Update train.py

train.py

@@ -10,31 +10,25 @@ from transformers import (
     AutoConfig
 )
 
-# ---- STEP 1: Load data ----
+# STEP 1: Define labels
 label_cols = [
     "Cog_present", "Aff_present", "Self_present",
-    "Motivation_present", "Attention_present", "OB_present", "Context_present"
+    "Motivation_present", "Attention_present", "OB_present", "Context_present",
+    "Social", "Physical", "Psych"
 ]
 
+# STEP 2: Load dataset
 df = pd.read_csv("/tmp/eemm_cleaned.csv")
 df_final = df[["clean_question"] + label_cols]
-
-# ---- STEP 2: Convert to Hugging Face dataset ----
 dataset = Dataset.from_pandas(df_final)
 
-# ---- STEP 3: Choose the best base model (changeable) ----
-base_model_name = "microsoft/deberta-v3-small"  # 🔥 Strong multi-label base model
-
-# ---- STEP 4: Tokenize ----
+# STEP 3: Choose model and tokenizer
+base_model_name = "microsoft/deberta-v3-small"
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
 
+# STEP 4: Tokenization
 def tokenize_and_format(example):
-    tokenized = tokenizer(
-        example["clean_question"],
-        padding="max_length",
-        truncation=True,
-        max_length=128
-    )
+    tokenized = tokenizer(example["clean_question"], padding="max_length", truncation=True, max_length=128)
     for label in label_cols:
         tokenized[label] = example[label]
     return tokenized
@@ -42,21 +36,20 @@ def tokenize_and_format(example):
 tokenized_dataset = dataset.map(tokenize_and_format)
 tokenized_dataset.set_format(type="torch", columns=["input_ids", "attention_mask"] + label_cols)
 
-# ---- STEP 5: Split into train/test ----
+# STEP 5: Train/test split
 train_test = tokenized_dataset.train_test_split(test_size=0.2)
 train_dataset = train_test["train"]
 eval_dataset = train_test["test"]
 
-# ---- STEP 6: Load model config and model ----
+# STEP 6: Model config and loading
 config = AutoConfig.from_pretrained(
     base_model_name,
     num_labels=len(label_cols),
     problem_type="multi_label_classification"
 )
-
 model = AutoModelForSequenceClassification.from_pretrained(base_model_name, config=config)
 
-# ---- STEP 7: Trainer setup ----
+# STEP 7: TrainingArguments and Trainer
 training_args = TrainingArguments(
     output_dir="./results",
     evaluation_strategy="epoch",
@@ -77,10 +70,10 @@ trainer = Trainer(
     eval_dataset=eval_dataset
 )
 
-# ---- STEP 8: Train ----
+# STEP 8: Train
 trainer.train()
 print("✅ Training complete.")
 
-# ---- STEP 9: Save model ----
+# STEP 9: Save
 trainer.save_model("./results")
 print("✅ Model saved to ./results")