Update backend/app/train.py

backend/app/train.py

@@ -2,10 +2,10 @@ import json
 import os
 from transformers import (
     AutoTokenizer,
-    AutoModelForCausalLM,
+    AutoModelForSequenceClassification,
     TrainingArguments,
     Trainer,
-    DataCollatorForLanguageModeling,
+    DataCollatorWithPadding,
 )
 import torch
 from datasets import Dataset
@@ -15,62 +15,48 @@ MODEL_NAME = "google/gemma-1.1-2b-it"
 DATA_PATH = "./backend/data/pregnancy_dataset.json"
 SAVE_PATH = "./backend/app/checkpoints"
 
-# -------- Load Dataset --------
-def load_dataset():
+# -------- Load and Preprocess Dataset --------
+def load_and_prepare_dataset():
     with open(DATA_PATH, "r") as f:
         data = json.load(f)
-    dataset = Dataset.from_list(data)
-    return dataset
 
-# -------- Tokenization --------
-def tokenize_function(example, tokenizer):
-    prompt = (
-        f"Age: {example['Age']}, SystolicBP: {example['SystolicBP']}, "
-        f"DiastolicBP: {example['DiastolicBP']}, BS: {example['BS']}, "
-        f"BodyTemp: {example['BodyTemp']}, HeartRate: {example['HeartRate']}. "
-        f"Predict the Risk Level."
-    )
-    completion = example["RiskLevel"]
-
-    # Map string labels to integer class indices
+    # Map risk levels to integer labels
     label_map = {"low risk": 0, "medium risk": 1, "high risk": 2}
-    label = label_map.get(completion.lower(), -1)  # -1 as fallback
 
-    full_prompt = f"<start_of_turn> {prompt} <end_of_turn>\n{completion} <end_of_turn>"
+    def preprocess(example):
+        prompt = (
+            f"Age: {example['Age']}, SystolicBP: {example['SystolicBP']}, "
+            f"DiastolicBP: {example['DiastolicBP']}, BS: {example['BS']}, "
+            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
+        return {"text": prompt, "label": label}
+
+    dataset = Dataset.from_list(data)
+    return dataset.map(preprocess)
 
-    tokenized = tokenizer(
-        full_prompt,
+# -------- Tokenization --------
+def tokenize_function(example, tokenizer):
+    return tokenizer(
+        example["text"],
         truncation=True,
-        padding="max_length",
+        padding=True,
         max_length=256,
     )
-    tokenized["labels"] = label  # ✅ Use numerical label
-
-    return tokenized
 
 # -------- Main Training Function --------
 def train():
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-
-    # Optional: Freeze all except final layers
-    for param in model.base_model.parameters():
-        param.requires_grad = False
-    for param in model.lm_head.parameters():
-        param.requires_grad = True
+    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=3)
 
-    # Load and tokenize
-    raw_dataset = load_dataset()
-    tokenized_dataset = raw_dataset.map(lambda x: tokenize_function(x, tokenizer))
+    dataset = load_and_prepare_dataset()
+    tokenized_dataset = dataset.map(lambda x: tokenize_function(x, tokenizer), batched=True)
 
-    # Data collator
-    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
-
-    # Training arguments
     training_args = TrainingArguments(
         output_dir=SAVE_PATH,
         num_train_epochs=3,
-        per_device_train_batch_size=2,
+        per_device_train_batch_size=4,
         save_steps=50,
         logging_steps=10,
         save_total_limit=1,
@@ -78,7 +64,8 @@ def train():
         report_to="none",
     )
 
-    # Trainer
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+
     trainer = Trainer(
         model=model,
         args=training_args,
@@ -94,3 +81,4 @@ def train():
 
 if __name__ == "__main__":
     train()
+