Update app_gpu.py

app_gpu.py

@@ -1,31 +1,41 @@
-# universal_lora_trainer_gradio_tabs.py
-import spaces
+# universal_lora_trainer_gradio.py
 import os
 import torch
 import gradio as gr
 import pandas as pd
+import numpy as np
 from pathlib import Path
 from torch.utils.data import Dataset, DataLoader
 from peft import LoraConfig, get_peft_model
 from accelerate import Accelerator
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import create_repo, upload_folder
-from tempfile import TemporaryDirectory
+from huggingface_hub import create_repo, upload_folder, hf_hub_download
+
+# transformers optional
+try:
+    from transformers import AutoTokenizer, AutoModelForCausalLM
+    TRANSFORMERS_AVAILABLE = True
+except Exception:
+    TRANSFORMERS_AVAILABLE = False
+
+# ---------------- Helpers ----------------
+def is_hub_repo_like(s):
+    return "/" in s and not Path(s).exists()
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+def download_from_hf(repo_id, filename, token=None):
+    token = token or os.environ.get("HF_TOKEN")
+    return hf_hub_download(repo_id=repo_id, filename=filename, repo_type="dataset", token=token)
 
 # ---------------- Dataset ----------------
 class MediaTextDataset(Dataset):
     def __init__(self, source, csv_name="dataset.csv", text_columns=None, max_records=None):
-        self.is_hub = "/" in source and not Path(source).exists()
+        self.is_hub = is_hub_repo_like(source)
         token = os.environ.get("HF_TOKEN")
         if self.is_hub:
-            from huggingface_hub import hf_hub_download
-            file_path = hf_hub_download(repo_id=source, filename=csv_name, repo_type="dataset", token=token)
+            file_path = download_from_hf(source, csv_name, token)
         else:
             file_path = Path(source) / csv_name
 
-        # fallback to parquet
+        # fallback to parquet if CSV missing
         if not Path(file_path).exists():
             alt = Path(str(file_path).replace(".csv", ".parquet"))
             if alt.exists():
@@ -36,6 +46,7 @@ class MediaTextDataset(Dataset):
         self.df = pd.read_parquet(file_path) if str(file_path).endswith(".parquet") else pd.read_csv(file_path)
         if max_records:
             self.df = self.df.head(max_records)
+
         self.text_columns = text_columns or ["short_prompt", "long_prompt"]
 
     def __len__(self):
@@ -43,18 +54,21 @@ class MediaTextDataset(Dataset):
 
     def __getitem__(self, i):
         rec = self.df.iloc[i]
-        return {"text": {col: rec[col] if col in rec else "" for col in self.text_columns}}
-
-# ---------------- Model helpers ----------------
-def load_pipeline(base_model, lora_repo=None, dtype=torch.float16):
-    tokenizer = AutoTokenizer.from_pretrained(base_model)
-    model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=dtype)
-    if lora_repo:
-        from peft import PeftModel
-        model = PeftModel.from_pretrained(model, lora_repo)
-    model.to(DEVICE)
-    model.eval()
-    return model, tokenizer
+        out = {"text": {}}
+        for col in self.text_columns:
+            out["text"][col] = rec[col] if col in rec else ""
+        return out
+
+# ---------------- Model loader ----------------
+def load_pipeline_auto(base_model, dtype=torch.float16):
+    if "gemma" in base_model.lower():
+        if not TRANSFORMERS_AVAILABLE:
+            raise RuntimeError("Transformers not installed for LLM support.")
+        tokenizer = AutoTokenizer.from_pretrained(base_model)
+        model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=dtype)
+        return {"model": model, "tokenizer": tokenizer}
+    else:
+        raise NotImplementedError("Only Gemma LLM supported in this script.")
 
 def find_target_modules(model):
     candidates = ["q_proj", "k_proj", "v_proj", "out_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
@@ -65,35 +79,57 @@ def find_target_modules(model):
     return targets
 
 def unwrap_batch(batch, short_col, long_col):
+    if isinstance(batch, (list, tuple)):
+        ex = batch[0]
+        if "text" in ex:
+            return ex
+        if "short" in ex and "long" in ex:
+            return {"text": {short_col: ex.get("short",""), long_col: ex.get("long","")}}
+        return {"text": ex}
+
     if isinstance(batch, dict):
+        first_elem = {}
+        is_batched = any(isinstance(v, (list, tuple, np.ndarray, torch.Tensor)) for v in batch.values())
+        if is_batched:
+            for k, v in batch.items():
+                try: first = v[0]
+                except Exception: first = v
+                first_elem[k] = first
+            if "text" in first_elem:
+                t = first_elem["text"]
+                if isinstance(t, (list, tuple)) and len(t) > 0:
+                    return {"text": t[0] if isinstance(t[0], dict) else {short_col: t[0], long_col: ""}}
+                if isinstance(t, dict): return {"text": t}
+                return {"text": {short_col: str(t), long_col: ""}}
+            if ("short" in first_elem and "long" in first_elem) or (short_col in first_elem and long_col in first_elem):
+                s = first_elem.get(short_col, first_elem.get("short", ""))
+                l = first_elem.get(long_col, first_elem.get("long", ""))
+                return {"text": {short_col: str(s), long_col: str(l)}}
+            return {"text": {short_col: str(first_elem)}}
+        if "text" in batch and isinstance(batch["text"], dict):
+            return {"text": batch["text"]}
         s = batch.get(short_col, batch.get("short", ""))
         l = batch.get(long_col, batch.get("long", ""))
         return {"text": {short_col: str(s), long_col: str(l)}}
-    if isinstance(batch, (list, tuple)):
-        ex = batch[0]
-        return unwrap_batch(ex, short_col, long_col)
     return {"text": {short_col: str(batch), long_col: ""}}
 
-# ---------------- Training ----------------
-import spaces
+# ---------------- LoRA Training ----------------
+from tempfile import TemporaryDirectory
+from accelerate import Accelerator
 
 @spaces.GPU(duration=110)
 def train_lora_stream(base_model, dataset_src, csv_name, text_cols,
                       epochs=1, lr=1e-4, r=8, alpha=16, batch_size=1,
                       num_workers=0, max_train_records=None, hf_repo_id=None):
 
-    if not hf_repo_id:
-        raise ValueError("❌ HF repo ID is required for upload.")
-    HF_TOKEN = os.environ.get("HF_TOKEN")
-    if not HF_TOKEN:
-        raise ValueError("❌ HF_TOKEN missing.")
-
-    dtype = torch.float16 if DEVICE == "cuda" else torch.float32
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float16 if device=="cuda" else torch.float32
     accelerator = Accelerator()
-    tokenizer = AutoTokenizer.from_pretrained(base_model)
-    model_obj = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=dtype)
-    model_obj.train()
+    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)
     lora_config = LoraConfig(r=r, lora_alpha=alpha, target_modules=target_modules, lora_dropout=0.0)
     lora_module = get_peft_model(model_obj, lora_config)
@@ -107,48 +143,81 @@ def train_lora_stream(base_model, dataset_src, csv_name, text_cols,
     step_counter = 0
     logs = []
 
-    yield "[INFO] Starting LoRA training...\n", 0.0
+    yield f"[INFO] Starting LoRA training on {device.upper()} (max {max_steps} steps)...\n", 0.0
 
     for ep in range(epochs):
+        yield f"[DEBUG] Epoch {ep+1}/{epochs}\n", step_counter / max_steps
         for batch in loader:
             if step_counter >= max_steps:
                 break
+
             ex = unwrap_batch(batch, text_cols[0], text_cols[1])
-            enc = tokenizer(ex["text"][text_cols[0]], text_pair=ex["text"][text_cols[1]],
-                            return_tensors="pt", padding="max_length", truncation=True, max_length=512)
-            enc = {k: v.to(accelerator.device) for k, v in enc.items()}
+            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)
+                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
-            logs.append(f"Step {step_counter}, Loss {loss.item():.4f}")
             yield "\n".join(logs[-10:]), step_counter / max_steps
+
         if step_counter >= max_steps:
             break
 
-    # Upload to HF
+    # ---------------- Upload to HF ----------------
+    HF_TOKEN = os.environ.get("HF_TOKEN")
+    if not hf_repo_id:
+        raise ValueError("❌ HF repo ID required for upload.")
+    if not HF_TOKEN:
+        raise ValueError("❌ HF_TOKEN missing.")
+
+    hf_repo_id = hf_repo_id.strip()
+    logs.append(f"[INFO] 🚀 Uploading LoRA to Hugging Face repo: {hf_repo_id}")
     create_repo(hf_repo_id, repo_type="model", exist_ok=True, token=HF_TOKEN)
+
     with TemporaryDirectory() as tmp_dir:
         lora_module.save_pretrained(tmp_dir)
         upload_folder(folder_path=tmp_dir, repo_id=hf_repo_id, repo_type="model", token=HF_TOKEN)
+
     link = f"https://huggingface.co/{hf_repo_id}"
     logs.append(f"[INFO] ✅ Uploaded successfully: {link}")
     yield "\n".join(logs), link
 
-# ---------------- Inference ----------------
-def generate_long_prompt(base_model, lora_repo, short_prompt, max_length=200):
-    model, tokenizer = load_pipeline(base_model, lora_repo=lora_repo)
-    input_ids = tokenizer(short_prompt, return_tensors="pt").input_ids.to(DEVICE)
+# ---------------- CPU Inference ----------------
+def generate_long_prompt_cpu(base_model, lora_repo, short_prompt, max_length=200):
+    device = torch.device("cpu")  # force CPU
+    pipe = load_pipeline_auto(base_model)
+    model = pipe["model"].to(device)
+    tokenizer = pipe["tokenizer"]
+
+    # Load LoRA adapter from HF
+    lora_module = get_peft_model(model, LoraConfig(
+        r=8, lora_alpha=16, target_modules=find_target_modules(model)
+    ))
+    lora_module.load_adapter(lora_repo, device=device, adapter_name="default")
+    lora_module.eval()
+
+    input_ids = tokenizer(short_prompt, return_tensors="pt").input_ids.to(device)
     with torch.no_grad():
-        outputs = model.generate(input_ids, max_length=max_length, do_sample=True, top_p=0.95, top_k=50)
+        outputs = lora_module.generate(input_ids, max_length=max_length, do_sample=True, top_p=0.95, top_k=50)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # ---------------- Gradio UI ----------------
@@ -156,25 +225,27 @@ def run_ui():
     with gr.Blocks() as demo:
         gr.Markdown("# 🌐 Universal Dynamic LoRA Trainer & Inference")
 
-        with gr.Tab("Training"):
-            base_model = gr.Textbox(label="Base model", value="google/gemma-3-4b-it")
-            dataset = gr.Textbox(label="Dataset folder or HF repo", value="rahul7star/prompt-enhancer-dataset-01")
-            csvname = gr.Textbox(label="CSV/Parquet file", value="train-00000-of-00001.csv")
-            short_col = gr.Textbox(label="Short prompt column", value="short_prompt")
-            long_col = gr.Textbox(label="Long prompt column", value="long_prompt")
-            repo = gr.Textbox(label="HF repo ID for LoRA upload (required)", value="rahul7star/gemma-3-270m-ccebc0")
-
-            batch_size = gr.Number(value=1, label="Batch size")
-            num_workers = gr.Number(value=0, label="DataLoader num_workers")
-            r = gr.Number(value=8, label="LoRA rank")
-            a = gr.Number(value=16, label="LoRA alpha")
-            ep = gr.Number(value=1, label="Epochs")
-            lr = gr.Number(value=1e-4, label="Learning rate")
-            max_records = gr.Number(value=1000, label="Max training records")
+        with gr.Tab("Train LoRA"):
+            with gr.Row():
+                base_model = gr.Textbox(label="Base model", value="google/gemma-3-4b-it")
+                dataset = gr.Textbox(label="Dataset folder or HF repo", value="rahul7star/prompt-enhancer-dataset-01")
+                csvname = gr.Textbox(label="CSV/Parquet file", value="train-00000-of-00001.csv")
+                short_col = gr.Textbox(label="Short prompt column", value="short_prompt")
+                long_col = gr.Textbox(label="Long prompt column", value="long_prompt")
+                repo = gr.Textbox(label="HF repo to upload LoRA", value="rahul7star/gemma-3-270m-ccebc0")
+
+            with gr.Row():
+                batch_size = gr.Number(value=1, label="Batch size")
+                num_workers = gr.Number(value=0, label="DataLoader num_workers")
+                r = gr.Number(value=8, label="LoRA rank")
+                a = gr.Number(value=16, label="LoRA alpha")
+                ep = gr.Number(value=1, label="Epochs")
+                lr = gr.Number(value=1e-4, label="Learning rate")
+                max_records = gr.Number(value=1000, label="Max training records")
 
             logs = gr.Textbox(label="Logs (streaming)", lines=25)
 
-            def launch_training(bm, ds, csv, sc, lc, batch, num_w, r_, a_, ep_, lr_, max_rec, repo_):
+            def launch_train(bm, ds, csv, sc, lc, batch, num_w, r_, a_, ep_, lr_, max_rec, repo_):
                 gen = train_lora_stream(
                     bm, ds, csv, [sc, lc],
                     epochs=int(ep_), lr=float(lr_), r=int(r_), alpha=int(a_),
@@ -185,23 +256,20 @@ def run_ui():
                     yield item
 
             btn = gr.Button("🚀 Start Training")
-            btn.click(fn=launch_training,
+            btn.click(fn=launch_train,
                       inputs=[base_model, dataset, csvname, short_col, long_col,
                               batch_size, num_workers, r, a, ep, lr, max_records, repo],
                       outputs=[logs],
                       queue=True)
 
-        with gr.Tab("Inference"):
+        with gr.Tab("Inference (CPU)"):
             inf_base_model = gr.Textbox(label="Base model", value="google/gemma-3-4b-it")
-            inf_lora_repo = gr.Textbox(label="LoRA HF repo ID", value="rahul7star/gemma-3-270m-ccebc0")
-            short_prompt = gr.Textbox(label="Short Prompt")
-            long_prompt_out = gr.Textbox(label="Generated Long Prompt", lines=5)
-
-            def run_inference(bm, lora_repo, sp):
-                return generate_long_prompt(bm, lora_repo, sp)
+            inf_lora_repo = gr.Textbox(label="LoRA HF repo", value="rahul7star/gemma-3-270m-ccebc0")
+            short_prompt = gr.Textbox(label="Short prompt")
+            long_prompt_out = gr.Textbox(label="Generated long prompt", lines=5)
 
             inf_btn = gr.Button("📝 Generate Long Prompt")
-            inf_btn.click(fn=run_inference,
+            inf_btn.click(fn=generate_long_prompt_cpu,
                           inputs=[inf_base_model, inf_lora_repo, short_prompt],
                           outputs=[long_prompt_out])