Update utils.py

utils.py

@@ -237,15 +237,20 @@ def convert_to_markdown(text):
     return markdown_text
 
 
-class State:
-    interrupted = False
+#Datasets encodieren - in train und val Sets
+class Dataset(torch.utils.data.Dataset):
+    def __init__(self, encodings, labels=None):
+        self.encodings = encodings
+        self.labels = labels
 
-    def interrupt(self):
-        self.interrupted = True
+    def __getitem__(self, idx):
+        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
+        if self.labels:
+            item["labels"] = torch.tensor(self.labels[idx])
+        return item
 
-    def recover(self):
-        self.interrupted = False
-shared_state = State()
+    def __len__(self):
+        return len(self.encodings["input_ids"])
 
 
 #######################################################