Update backend/app/train.py

backend/app/train.py

@@ -30,7 +30,9 @@ def load_and_prepare_dataset():
             f"BodyTemp: {example['BodyTemp']}, HeartRate: {example['HeartRate']}. "
             f"Predict the Risk Level."
         )
-        label = label_map.get(example["RiskLevel"].lower(), 0)  # Default to 0 if unknown
+        # Ensure consistent and safe label mapping
+        label_str = str(example.get("RiskLevel", "")).lower()
+        label = label_map.get(label_str, 0)
         return {"text": prompt, "label": label}
 
     dataset = Dataset.from_list(data)
@@ -44,7 +46,7 @@ def tokenize_function(example, tokenizer):
         padding=True,
         max_length=256,
     )
-    tokens["label"] = example["label"]  # ✅ Keep label after tokenization
+    tokens["label"] = example["label"]
     return tokens
 
 # -------- Main Training Function --------
@@ -53,8 +55,17 @@ def train():
     model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=3)
 
     dataset = load_and_prepare_dataset()
+
+    # Tokenize dataset
     tokenized_dataset = dataset.map(lambda x: tokenize_function(x, tokenizer), batched=False)
-    tokenized_dataset = tokenized_dataset.remove_columns(["text"])
+
+    # Remove any non-tensor-compatible fields
+    tokenized_dataset = tokenized_dataset.remove_columns(
+        [col for col in tokenized_dataset.column_names if col not in ["input_ids", "attention_mask", "label"]]
+    )
+
+    # Optional sanity check
+    print("🔎 Sample tokenized example:", tokenized_dataset[0])
 
     training_args = TrainingArguments(
         output_dir=SAVE_PATH,