Update app.py

app.py

@@ -9,11 +9,12 @@ from collections import deque
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModel
+from datasets import load_dataset
 
-ACTIONS = ["TRIP", "GITHUB", "MAIL"]
+ACTIONS = ["TRIP", "GITHUB", "MAIL", "CALENDAR"]
 NUM_ACTIONS = len(ACTIONS)
 DATASET_PATH = "iteratehack/code19-dataset"
-
+HF_SPLIT = "train"
 # Confidence threshold - below this returns NONE
 CONFIDENCE_THRESHOLD = 0.6
 
@@ -230,7 +231,6 @@ class RLAgent:
 
 
 def load_dataset():
-    """Load dataset from Hugging Face Datasets."""
     dataset = load_dataset(HF_DATASET, split=HF_SPLIT)
 
     data = []