Update train.py

train.py

@@ -3,10 +3,8 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trai
 from peft import get_peft_model, LoraConfig, TaskType
 import os
 
-# Load SST2 dataset from GLUE (binary sentiment classification)
+# Load SST-2 dataset (sentiment classification) and take a small subset for fast training
 dataset = load_dataset("glue", "sst2")
-
-# Use a small subset to stay within 25-minute budget
 small_train = dataset["train"].select(range(500))
 
 tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
@@ -16,15 +14,23 @@ def tokenize_fn(batch):
 
 tokenized_train = small_train.map(tokenize_fn, batched=True)
 
-# Load model and apply LoRA
+# Load model
 model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)
-peft_config = LoraConfig(task_type=TaskType.SEQ_CLS, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1)
+
+# Apply PEFT with LoRA — FIXED: target_modules is now set for DistilBERT
+peft_config = LoraConfig(
+    task_type=TaskType.SEQ_CLS,
+    inference_mode=False,
+    r=8,
+    lora_alpha=32,
+    lora_dropout=0.1,
+    target_modules=["q_lin", "v_lin"]  # Required for DistilBERT
+)
 model = get_peft_model(model, peft_config)
 
-# Hugging Face token from environment or manually
-hf_token = os.environ.get("HF_TOKEN") or "hf_xxx"  # replace with real token or set in Space secrets
+# Hugging Face token (set as a Secret in Space settings)
+hf_token = os.environ.get("HF_TOKEN") or "hf_xxx"  # Replace if needed
 
-# Training arguments
 training_args = TrainingArguments(
     output_dir="results",
     per_device_train_batch_size=8,
@@ -34,7 +40,7 @@ training_args = TrainingArguments(
     save_strategy="epoch",
     push_to_hub=True,
     hub_model_id="NightPrince/peft-distilbert-sst2",
-    hub_token=hf_token,
+    hub_token=hf_token
 )
 
 trainer = Trainer(