From 747c845fb1464f03269ea84409908d03487adb68 Mon Sep 17 00:00:00 2001
From: sedherthe <siddhartha.soma1@gmail.com>
Date: Sat, 21 Feb 2026 01:46:44 +0000
Subject: [PATCH] Add updated dataset for training llm.

---
 dataset.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)
 create mode 100644 dataset.py

diff --git a/dataset.py b/dataset.py
new file mode 100644
index 0000000..5e3eb17
--- /dev/null
+++ b/dataset.py
@@ -0,0 +1,17 @@
+import json
+from torch.utils.data import Dataset
+
+
+class AudioDataset(Dataset):
+    def __init__(self, path):
+        with open(path, encoding='utf-8') as f:
+            self.dataset = json.load(f)
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        text, audio, audio_path = self.dataset[idx]
+        # Format: [STOP][TEXT]<text prompt>[START]<audio tokens>[STOP]
+        res = f"[TEXT]{text}[START]{''.join(list(map(lambda x: f'[{x}]', audio)))}[STOP]"
+        return res