Update app_gpu.py

app_gpu.py

@@ -139,113 +139,124 @@ import spaces
 import torch
 from huggingface_hub import create_repo, upload_folder
 
-@spaces.GPU(duration=100)
-def train_lora_model(
-    base_model,
-    dataset,
-    csvname,
-    short_col,
-    long_col,
-    out,
-    repo,
-    batch_size,
-    num_workers,
-    r,
-    a,
-    ep,
-    lr,
-    max_records
-):
-    import torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
-    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
-    from datasets import load_dataset
-    import os
-    from huggingface_hub import HfApi, upload_folder
-
-    print(f"[INFO] Loading base model: {base_model}")
+
+@spaces.GPU(duration=110)
+def train_lora_stream(base_model, dataset_src, csv_name, text_cols, output_dir,
+                      epochs=1, lr=1e-4, r=8, alpha=16, batch_size=1, num_workers=0,
+                      max_train_records=None, repo_id=None):
+    """LoRA training loop with GPU for compute, CPU for upload, capped at 150 steps."""
+
+    # --- Device setup ---
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32)
-    tokenizer = AutoTokenizer.from_pretrained(base_model)
-    tokenizer.pad_token = tokenizer.eos_token
-
-    # LoRA setup
-    print("[INFO] Setting up LoRA configuration...")
-    lora_config = LoraConfig(
-        r=r,
-        lora_alpha=a,
-        target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
-        lora_dropout=0.05,
-        bias="none",
-        task_type="CAUSAL_LM"
-    )
-    model = prepare_model_for_kbit_training(model)
-    model = get_peft_model(model, lora_config)
-    model.print_trainable_parameters()
-
-    print(f"[INFO] Loading dataset from: {dataset}")
-    ds = load_dataset(dataset)
-    df = ds["train"].to_pandas()
-    print(f"[DEBUG] Loaded dataset: {dataset}, columns: {df.columns.tolist()}")
-    print("[DEBUG] Sample rows:\n", df.head(3))
-
-    df = df[[short_col, long_col]].dropna().head(max_records)
-    train_data = list(zip(df[short_col], df[long_col]))
-
-    print(f"[INFO] Tokenizing {len(train_data)} records...")
-    def tokenize(examples):
-        inputs = [f"Short: {s}\nLong: {l}" for s, l in examples]
-        model_inputs = tokenizer(inputs, truncation=True, padding="max_length", max_length=512)
-        model_inputs["labels"] = model_inputs["input_ids"].copy()
-        return model_inputs
-
-    tokenized_data = [tokenize([f"{s}\n{l}"]) for s, l in train_data]
-    print(f"[INFO] Tokenized {len(tokenized_data)} samples")
-
-    # Trainer setup
-    training_args = TrainingArguments(
-        output_dir=out,
-        num_train_epochs=ep,
-        per_device_train_batch_size=batch_size,
-        learning_rate=lr,
-        logging_dir=os.path.join(out, "logs"),
-        logging_steps=10,
-        save_strategy="no",
-        report_to="none",
-        dataloader_num_workers=num_workers,
-        max_steps=200,  # ✅ Limit steps to 200 to avoid timeout
-    )
-
-    print("[INFO] Starting training loop (max 200 steps)...")
-    dummy_dataset = [{"input_ids": torch.tensor(d["input_ids"]), "labels": torch.tensor(d["labels"])} for d in tokenized_data]
-
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=dummy_dataset
-    )
-
-    trainer.train()
-    print("[✅] Training completed!")
-
-    # Save adapter and tokenizer locally first
-    model.save_pretrained(out)
-    tokenizer.save_pretrained(out)
-    print(f"[INFO] Model saved locally at: {out}")
-
-    # ✅ Upload happens in CPU mode only (after freeing GPU memory)
-    if repo:
-        print("[INFO] Switching to CPU for model upload...")
-        del model
+    gpu_name = torch.cuda.get_device_name(0) if device == "cuda" else "CPU"
+    print(f"[INFO] 🚀 Using device: {device.upper()} ({gpu_name})")
+
+    # Adjust precision / batch based on VRAM
+    if device == "cuda":
+        vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+        print(f"[INFO] VRAM: {vram:.2f} GB")
+        dtype = torch.bfloat16 if "A100" in gpu_name or vram > 20 else torch.float16
+        if vram < 10:
+            batch_size = max(1, batch_size // 2)
+            print(f"[WARN] Low VRAM, using batch_size={batch_size}")
+    else:
+        dtype = torch.float32
+
+    # --- Model & tokenizer ---
+    accelerator = Accelerator()
+    pipe = load_pipeline_auto(base_model, dtype=dtype)
+    model_obj = pipe["model"]
+    tokenizer = pipe["tokenizer"]
+
+    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)
+
+    # --- Dataset ---
+    dataset = MediaTextDataset(dataset_src, csv_name, text_columns=text_cols, max_records=max_train_records)
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+    optimizer = torch.optim.AdamW(lora_module.parameters(), lr=lr)
+    lora_module, optimizer, loader = accelerator.prepare(lora_module, optimizer, loader)
+
+    # --- Limit steps to 150 ---
+    max_steps = 150
+    total_steps = min(max_steps, max(1, epochs * len(loader)))
+    step_counter = 0
+    logs = []
+
+    yield f"[INFO] Starting LoRA training on {gpu_name} (max {max_steps} steps)...\n", 0.0
+
+    # --- Training Loop ---
+    for ep in range(epochs):
+        yield f"[DEBUG] Epoch {ep+1}/{epochs}\n", step_counter / total_steps
+        for i, batch in enumerate(loader):
+            if step_counter >= max_steps:
+                break
+
+            ex = unwrap_batch(batch, text_cols[0], text_cols[1])
+            texts = ex.get("text", {})
+            short_text = str(texts.get(text_cols[0], "") or "")
+            long_text = str(texts.get(text_cols[1], "") or "")
+
+            enc = tokenizer(
+                short_text,
+                text_pair=long_text,
+                return_tensors="pt",
+                padding="max_length",
+                truncation=True,
+                max_length=512,
+            )
+            enc = {k: v.to(accelerator.device) for k, v in enc.items()}
+            enc["labels"] = enc["input_ids"].clone()
+
+            outputs = lora_module(**enc)
+            loss = getattr(outputs, "loss", None)
+            if loss is None:
+                logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
+                loss = torch.nn.functional.cross_entropy(
+                    logits.view(-1, logits.size(-1)),
+                    enc["labels"].view(-1),
+                    ignore_index=tokenizer.pad_token_id
+                )
+
+            optimizer.zero_grad()
+            accelerator.backward(loss)
+            optimizer.step()
+
+            logs.append(f"[DEBUG] Step {step_counter}, Loss: {loss.item():.6f}")
+            step_counter += 1
+            yield "\n".join(logs[-10:]), step_counter / total_steps
+
+        if step_counter >= max_steps:
+            break
+
+    # --- Save LoRA ---
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+    lora_module.save_pretrained(output_dir)
+    yield f"[INFO] ✅ LoRA saved to {output_dir}\n", 0.95
+
+    # --- Free GPU before upload ---
+    if torch.cuda.is_available():
+        yield "[INFO] Releasing GPU memory before upload...\n", 0.96
+        del lora_module
         torch.cuda.empty_cache()
+        torch.cuda.synchronize()
 
-        api = HfApi()
-        print(f"[INFO] Uploading model to Hugging Face repo: {repo}")
-        upload_folder(folder_path=out, repo_id=repo, repo_type="model", commit_message="Upload trained LoRA adapter")
-        print(f"[✅] Model uploaded successfully to {repo}")
-
-    return f"✅ Training done and uploaded to {repo if repo else 'local directory only.'}"
+    # --- Upload to HF (CPU mode only) ---
+    repo_id = repo_id or os.environ.get("HF_UPLOAD_REPO")
+    token = os.environ.get("HF_TOKEN")
 
+    if repo_id and token:
+        yield f"[INFO] Uploading adapter to {repo_id} (CPU mode)...\n", 0.97
+        try:
+            create_repo(repo_id, repo_type="model", exist_ok=True, token=token)
+            upload_folder(folder_path=output_dir, repo_id=repo_id, repo_type="model", token=token)
+            yield f"[INFO] ✅ Uploaded successfully: https://huggingface.co/{repo_id}\n", 1.0
+        except Exception as e:
+            yield f"[ERROR] Upload failed: {e}\n", 1.0
+    else:
+        yield f"[INFO] Skipping upload — repo_id or token not provided.\n", 1.0
 
 
 def upload_adapter(local, repo_id):
@@ -282,7 +293,7 @@ def run_ui():
         logs = gr.Textbox(label="Logs (streaming)", lines=25)
 
         def launch(bm, ds, csv, sc, lc, out_dir, batch, num_w, r_, a_, ep_, lr_, max_rec, repo_):
-            gen = train_lora_model(
+            gen = train_lora_stream(
                 bm, ds, csv, [sc, lc], out_dir,
                 epochs=int(ep_), lr=float(lr_), r=int(r_), alpha=int(a_),
                 batch_size=int(batch), num_workers=int(num_w),