Blame for app.py

import math
import os
import re
import tempfile
import time

import gradio as gr
import spaces
import torch
from datasets import load_dataset
from huggingface_hub import HfApi
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer

BASE_MODELS = [
    "HuggingFaceTB/SmolLM2-135M-Instruct",
    "HuggingFaceTB/SmolLM2-360M-Instruct",
    "HuggingFaceTB/SmolLM2-1.7B-Instruct",
]
DEFAULT_DATASET = "x"
DEFAULT_TARGET = "x"  # gilt nur für den Besitzer
OWNER = "x"
MAX_LEN = 384
BATCH_SIZE = 8


def clean_id(s: str) -> str:
    """Akzeptiert URL oder ID: https://huggingface.co/datasets/user/name -> user/name"""
    s = (s or "").strip().rstrip("/")
    s = re.sub(r"^https?://huggingface\.co/", "", s)
    s = re.sub(r"^(datasets|models)/", "", s)
    s = re.sub(r"/(tree|blob|viewer)/.*$", "", s)
    return s


def build_texts(ds, tok, text_col):
    cols = ds.column_names
    texts = []
    for row in ds:
        if "messages" in cols:
            t = tok.apply_chat_template(row["messages"], tokenize=False)
        elif text_col and text_col in cols:
            t = str(row[text_col])
        elif "text" in cols:
            t = str(row["text"])
        else:
            q = next((row[c] for c in ("instruction", "prompt", "question", "input") if c in cols and row[c]), None)
            a = next((row[c] for c in ("output", "response", "answer", "completion") if c in cols and row[c]), None)
            if q is None or a is None:
                raise gr.Error(
                    f"Spalten nicht erkannt: {cols}. Nutze 'messages', 'text' oder instruction/output "
                    "(oder gib den Spaltennamen an)."
                )
            t = tok.apply_chat_template(
                [{"role": "user", "content": str(q)}, {"role": "assistant", "content": str(a)}],
                tokenize=False,
            )
        texts.append(t)
    return texts


def gpu_duration(model, input_ids, attention_mask, seconds, epochs, lr, out_dir):
    return int(seconds) + 20


@spaces.GPU(duration=gpu_duration)
def train_gpu(model, input_ids, attention_mask, seconds, epochs, lr, out_dir):
    model.to("cuda")
    model.train()
    params = [p for p in model.parameters() if p.requires_grad]
    opt = torch.optim.AdamW(params, lr=lr)
    n = input_ids.size(0)
    start = time.time()
    losses, stop = [], False
    for _ in range(int(epochs)):
        perm = torch.randperm(n)
        for i in range(0, n, BATCH_SIZE):
            if time.time() - start > seconds:
                stop = True
                break
            idx = perm[i : i + BATCH_SIZE]
            ids = input_ids[idx].to("cuda")
            am = attention_mask[idx].to("cuda")
            labels = ids.clone()
            labels[am == 0] = -100
            loss = model(input_ids=ids, attention_mask=am, labels=labels).loss
            loss.backward()
            opt.step()
            opt.zero_grad()
            losses.append(loss.item())
        if stop:
            break
    model.save_pretrained(out_dir)
    return losses, time.time() - start, stop


def run(
    base, dataset_in, split, text_col, max_examples, epochs, seconds, lr, rank,
    target_repo, manual_token, profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None,
    progress=gr.Progress(),
):
    token = (manual_token or "").strip() or (oauth_token.token if oauth_token else None)
    if not token:
        raise gr.Error("Bitte oben mit Hugging Face einloggen (oder Token eingeben).")
    api = HfApi(token=token)
    user = profile.username if profile else api.whoami()["name"]
    dataset_id = clean_id(dataset_in)
    if not dataset_id:
        raise gr.Error("Dataset-URL oder -ID fehlt.")
    if clean_id(target_repo):
        repo = clean_id(target_repo)
    elif user.lower() == OWNER.lower():
        repo = DEFAULT_TARGET
    else:
        repo = f"{user}/{base.split('/')[-1]}-ft"

    yield "Lade Tokenizer & Dataset (CPU, zählt nicht zur GPU-Zeit) ..."
    tok = AutoTokenizer.from_pretrained(base)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    ds = load_dataset(dataset_id, split=f"{split}[:{int(max_examples)}]", token=token)
    texts = build_texts(ds, tok, text_col)
    enc = tok(texts, truncation=True, max_length=MAX_LEN, padding="max_length", return_tensors="pt")

    yield f"{len(texts)} Beispiele bereit. Lade Modell + LoRA (CPU) ..."
    model = AutoModelForCausalLM.from_pretrained(base, torch_dtype=torch.bfloat16)
    cfg = LoraConfig(
        r=int(rank), lora_alpha=int(rank) * 2, lora_dropout=0.05,
        target_modules="all-linear", task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, cfg)

    out_dir = tempfile.mkdtemp()
    yield f"Training startet (max. {int(seconds)} s GPU-Zeit) ..."
    losses, took, stopped = train_gpu(
        model, enc["input_ids"], enc["attention_mask"], float(seconds), int(epochs), float(lr), out_dir
    )
    if not losses:
        raise gr.Error("Kein Trainingsschritt gelaufen (GPU-Zeit zu kurz?).")
    first = sum(losses[:5]) / len(losses[:5])
    last = sum(losses[-5:]) / len(losses[-5:])

    yield f"Fertig trainiert ({len(losses)} Steps, {took:.0f} s). Lade hoch nach {repo} ..."
    api.create_repo(repo, exist_ok=True, repo_type="model")
    with open(os.path.join(out_dir, "README.md"), "w") as f:
        f.write(
            f"---\nbase_model: {base}\nlibrary_name: peft\ntags:\n- lora\n- zerogpu-trainer\n---\n"
            f"LoRA-Adapter auf `{base}`, trainiert mit `{dataset_id}`.\n"
        )
    api.upload_folder(folder_path=out_dir, repo_id=repo, repo_type="model")
    note = " (Zeitlimit erreicht, Training vorzeitig beendet)" if stopped else ""
    yield (
        f"Fertig{note}\n"
        f"Steps: {len(losses)} | Loss: {first:.3f} → {last:.3f}\n"
        f"Adapter: https://huggingface.co/{repo}"
    )


def load_repos(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None):
    """Listet die Datasets und Modelle des eingeloggten Nutzers für die Auswahl."""
    if not oauth_token or not profile:
        return gr.update(), gr.update(), "Noch nicht eingeloggt."
    try:
        api = HfApi(token=oauth_token.token)
        ds = sorted(d.id for d in api.list_datasets(author=profile.username, limit=200))
        md = sorted(m.id for m in api.list_models(author=profile.username, limit=200))
    except Exception as e:  # noqa: BLE001
        return gr.update(), gr.update(), f"Konnte Repos nicht laden: {e}"
    ds_choices = ds + ([DEFAULT_DATASET] if DEFAULT_DATASET not in ds else [])
    return (
        gr.update(choices=ds_choices),
        gr.update(choices=md),
        f"Eingeloggt als **{profile.username}** – {len(ds)} Datasets, {len(md)} Modelle gefunden.",
    )


with gr.Blocks(title="ZeroGPU Mini-Trainer") as demo:
    gr.Markdown(
        "# ⚡ ZeroGPU Mini-Trainer\n"
        "Fine-Tune ein kleines SmolLM per LoRA in wenigen Minuten – mit deiner eigenen ZeroGPU-Tagesquote."
    )
    gr.LoginButton()
    with gr.Row():
        refresh = gr.Button("🔄 Meine Datasets & Modelle laden")
        repo_info = gr.Markdown("Logge dich ein, dann erscheinen deine Repos in den Auswahlfeldern.")
    with gr.Row():
        base = gr.Dropdown(BASE_MODELS, value=BASE_MODELS[1], label="Basismodell", allow_custom_value=True)
        dataset_in = gr.Dropdown(
            choices=[DEFAULT_DATASET], value=DEFAULT_DATASET, allow_custom_value=True,
            label="Dataset (auswählen oder URL / user/name einfügen)",
        )
    with gr.Row():
        split = gr.Textbox(value="train", label="Split")
        text_col = gr.Textbox(label="Textspalte (optional)", placeholder="leer = automatisch")
        max_examples = gr.Slider(50, 5000, value=1000, step=50, label="Max. Beispiele")
    with gr.Row():
        epochs = gr.Slider(1, 5, value=2, step=1, label="Epochen")
        seconds = gr.Slider(30, 240, value=150, step=10, label="Max. GPU-Sekunden")
        lr = gr.Number(value=2e-4, label="Learning Rate")
        rank = gr.Slider(4, 64, value=16, step=4, label="LoRA Rank")
    target_repo = gr.Dropdown(
        choices=[], value=None, allow_custom_value=True,
        label="Ziel-Modell-Repo (auswählen oder eintippen; leer = automatisch)",
    )
    manual_token = gr.Textbox(label="Optional: HF-Token (nur falls Login nicht geht)", type="password")
    btn = gr.Button("Training starten", variant="primary")
    out = gr.Textbox(label="Status", lines=6)
    btn.click(
        run,
        [base, dataset_in, split, text_col, max_examples, epochs, seconds, lr, rank, target_repo, manual_token],
        out,
    )

    refresh.click(load_repos, None, [dataset_in, target_repo, repo_info])
    demo.load(load_repos, None, [dataset_in, target_repo, repo_info])

demo.queue().launch()