# universal_lora_trainer_accelerate_singlefile_dynamic.py """ Universal Dynamic LoRA Trainer (Accelerate + PEFT + Gradio) - Gemma LLM default - Robust batch handling (fixes KeyError: 0) - Streams logs to Gradio (includes progress %) - Supports CSV/Parquet HuggingFace or local datasets """ 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 huggingface_hub import hf_hub_download, create_repo, upload_folder # transformers optional try: from transformers import AutoTokenizer, AutoModelForCausalLM TRANSFORMERS_AVAILABLE = True except Exception: TRANSFORMERS_AVAILABLE = False DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # ---------------- Helpers ---------------- def is_hub_repo_like(s): return "/" in s and not Path(s).exists() 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 = is_hub_repo_like(source) token = os.environ.get("HF_TOKEN") if self.is_hub: file_path = download_from_hf(source, csv_name, token) else: file_path = Path(source) / csv_name # fallback to parquet if CSV missing if not Path(file_path).exists(): alt = Path(str(file_path).replace(".csv", ".parquet")) if alt.exists(): file_path = alt else: raise FileNotFoundError(f"Dataset file not found: {file_path}") 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"] print(f"[DEBUG] Loaded dataset: {file_path}, columns: {list(self.df.columns)}") print(f"[DEBUG] Sample rows:\n{self.df.head(3)}") def __len__(self): return len(self.df) def __getitem__(self, i): rec = self.df.iloc[i] 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.") print(f"[INFO] Using Gemma LLM for {base_model}") 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"] names = [n for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)] targets = [n.split(".")[-1] for n in names if any(c in n for c in candidates)] if not targets: targets = [n.split(".")[-1] for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)] print(f"[WARNING] No standard attention modules found, using Linear layers for LoRA.") else: print(f"[INFO] LoRA target modules detected: {targets[:40]}{'...' if len(targets)>40 else ''}") return targets # ---------------- Batch unwrapping ---------------- 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)}} return {"text": {short_col: str(batch), long_col: ""}} # ---------------- Training (forward + backward + logs) ---------------- 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): accelerator = Accelerator() pipe = load_pipeline_auto(base_model) 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 = 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) total_steps = max(1, epochs * len(loader)) step_counter = 0 logs = [] yield "[DEBUG] Starting training loop...\n", 0.0 for ep in range(epochs): yield f"[DEBUG] Epoch {ep+1}/{epochs}\n", step_counter / total_steps for i, batch in enumerate(loader): 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 "") # --- 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 ) enc = {k: v.to(accelerator.device) for k, v in enc.items()} enc["labels"] = enc["input_ids"].clone() # --- Forward pass --- outputs = lora_module(**enc) forward_loss = getattr(outputs, "loss", None) if forward_loss is None: logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] forward_loss = torch.nn.functional.cross_entropy( logits.view(-1, logits.size(-1)), enc["labels"].view(-1), ignore_index=tokenizer.pad_token_id ) logs.append(f"[DEBUG] Step {step_counter}, forward_loss: {forward_loss.item():.6f}") optimizer.zero_grad() accelerator.backward(forward_loss) optimizer.step() step_counter += 1 yield "\n".join(logs[-10:]), step_counter / total_steps Path(output_dir).mkdir(parents=True, exist_ok=True) lora_module.save_pretrained(output_dir) yield f"[INFO] ✅ LoRA saved to {output_dir}\n", 1.0 def upload_adapter(local, repo_id): token = os.environ.get("HF_TOKEN") if not token: raise RuntimeError("HF_TOKEN missing") create_repo(repo_id, exist_ok=True) upload_folder(local, repo_id=repo_id, repo_type="model", token=token) return f"https://huggingface.co/{repo_id}" # ---------------- Gradio UI ---------------- def run_ui(): with gr.Blocks() as demo: gr.Markdown("# 🌐 Universal Dynamic LoRA Trainer (Gemma LLM)") 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") out = gr.Textbox(label="Output dir", value="./adapter_out") repo = gr.Textbox(label="Upload HF repo (optional)", 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(bm, ds, csv, sc, lc, out_dir, batch, num_w, r_, a_, ep_, lr_, max_rec, repo_): 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), max_train_records=int(max_rec) ) for item in gen: if isinstance(item, tuple): text = item[0] else: text = item yield text if repo_: link = upload_adapter(out_dir, repo_) yield f"[INFO] Uploaded to {link}\n" btn = gr.Button("🚀 Start Training") btn.click(fn=launch, inputs=[base_model, dataset, csvname, short_col, long_col, out, batch_size, num_workers, r, a, ep, lr, max_records, repo], outputs=[logs], queue=True) return demo if __name__ == "__main__": run_ui().launch(server_name="0.0.0.0", server_port=7860, share=True)