Create Dataset.py

Dataset.py

@@ -0,0 +1,31 @@
+import pandas as pd
+from transformers import AutoTokenizer
+
+class MyDataset:
+    def __init__(self, data_file, tokenizer):
+        self.data = pd.read_csv(data_file)
+        self.tokenizer = tokenizer
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        text = self.data.iloc[idx, 0]
+        agents = self.data.iloc[idx, 1]
+        actions = self.data.iloc[idx, 2]
+
+        encoding = self.tokenizer.encode_plus(
+            text,
+            max_length=512,
+            padding='max_length',
+            truncation=True,
+            return_attention_mask=True,
+            return_tensors='pt'
+        )
+
+        return {
+            'input_ids': encoding['input_ids'].flatten(),
+            'attention_mask': encoding['attention_mask'].flatten(),
+            'labels_agents': torch.tensor(agents),
+            'labels_actions': torch.tensor(actions)
+        }