Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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() |