Update app.py

app.py

@@ -140,9 +140,7 @@ def train_lora_stream(base_model, dataset_src, csv_name, text_cols, output_dir,
     model_obj = pipe["model"]
     tokenizer = pipe["tokenizer"]
 
-    # Ensure model is in train mode
     model_obj.train()
-
     target_modules = find_target_modules(model_obj)
     lcfg = LoraConfig(r=r, lora_alpha=alpha, target_modules=target_modules, lora_dropout=0.0)
     lora_module = get_peft_model(model_obj, lcfg)
@@ -167,32 +165,31 @@ def train_lora_stream(base_model, dataset_src, csv_name, text_cols, output_dir,
             short_text = str(texts.get(text_cols[0], "") or "")
             long_text = str(texts.get(text_cols[1], "") or "")
 
-            # Encode both short and long as supervised pairs
-            inputs = tokenizer(short_text, return_tensors="pt", truncation=True, padding=True, max_length=1024)
-            labels = tokenizer(long_text, return_tensors="pt", truncation=True, padding=True, max_length=1024)
+            # --- FIX: Tokenize as text pair to align sequence lengths ---
+            enc = tokenizer(
+                short_text,
+                text_pair=long_text,
+                return_tensors="pt",
+                padding="max_length",
+                truncation=True,
+                max_length=512,  # enforce same length for both
+            )
 
-            inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
-            labels = labels["input_ids"].to(DEVICE)
+            enc = {k: v.to(accelerator.device) for k, v in enc.items()}
+            enc["labels"] = enc["input_ids"].clone()
 
-            # Forward pass with labels — ensures gradient flow
-            outputs = lora_module(**inputs, labels=labels)
+            # --- Forward pass ---
+            outputs = lora_module(**enc)
 
-            # Handle loss properly
             forward_loss = getattr(outputs, "loss", None)
             if forward_loss is None:
-                # Fallback to MSE loss between logits and labels
                 logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
-                forward_loss = torch.nn.functional.mse_loss(
-                    logits.float(), torch.nn.functional.one_hot(labels, num_classes=logits.size(-1)).float()
+                forward_loss = torch.nn.functional.cross_entropy(
+                    logits.view(-1, logits.size(-1)), enc["labels"].view(-1), ignore_index=tokenizer.pad_token_id
                 )
 
-            # Ensure loss requires grad
-            if not forward_loss.requires_grad:
-                forward_loss = forward_loss.clone().detach().requires_grad_(True)
-
             logs.append(f"[DEBUG] Step {step_counter}, forward_loss: {forward_loss.item():.6f}")
 
-            # Backprop
             optimizer.zero_grad()
             accelerator.backward(forward_loss)
             optimizer.step()