app.py

# 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)