Update app.py

app.py

@@ -12,7 +12,7 @@ from transformers import AutoTokenizer, AutoModel
 
 ACTIONS = ["TRIP", "GITHUB", "MAIL"]
 NUM_ACTIONS = len(ACTIONS)
-DATASET_PATH = os.path.join(os.path.dirname(__file__), "dataset.jsonl")
+DATASET_PATH = "iteratehack/code19-dataset"
 
 # Confidence threshold - below this returns NONE
 CONFIDENCE_THRESHOLD = 0.6
@@ -230,16 +230,15 @@ class RLAgent:
 
 
 def load_dataset():
-    """Load and parse the dataset."""
-    data = []
+    """Load dataset from Hugging Face Datasets."""
+    dataset = load_dataset(HF_DATASET, split=HF_SPLIT)
 
-    with open(DATASET_PATH, "r") as f:
-        for line in f:
-            item = json.loads(line)
-            user_msg = item["messages"][1]["content"]
-            label = item["messages"][2]["content"]
-            if label in ACTIONS:
-                data.append((user_msg, ACTIONS.index(label)))
+    data = []
+    for item in dataset:
+        user_msg = item["messages"][1]["content"]
+        label = item["messages"][2]["content"]
+        if label in ACTIONS:
+            data.append((user_msg, ACTIONS.index(label)))
 
     random.shuffle(data)
     return data