Create app.py

app.py

@@ -0,0 +1,776 @@
+"""
+╔══════════════════════════════════════════════════════════════╗
+║           🧪 Fine-Tuning Studio — HuggingFace Space          ║
+║  Suporta: CPU / CPU Upgrade / T4 / A10G / A100               ║
+║  Modos:   LoRA, QLoRA, Full Fine-Tuning                      ║
+║  Pós:     Chat embutido + Download dos pesos                  ║
+╚══════════════════════════════════════════════════════════════╝
+"""
+
+import os, gc, json, math, shutil, threading, time, logging
+from pathlib import Path
+from typing import Optional
+
+import gradio as gr
+import torch
+import pandas as pd
+from datasets import load_dataset, Dataset
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling,
+    BitsAndBytesConfig,
+    GenerationConfig,
+    TrainerCallback,
+)
+
+# ── PEFT é opcional; detectado em tempo de execução ───────────
+try:
+    from peft import (
+        LoraConfig,
+        get_peft_model,
+        prepare_model_for_kbit_training,
+        PeftModel,
+        TaskType,
+    )
+    PEFT_AVAILABLE = True
+except ImportError:
+    PEFT_AVAILABLE = False
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ──────────────────────────────────────────────────────────────
+#  HARDWARE DETECTION
+# ──────────────────────────────────────────────────────────────
+
+def detect_hardware() -> dict:
+    info = {"device": "cpu", "vram_gb": 0, "gpu_name": "N/A", "bf16": False}
+    if torch.cuda.is_available():
+        info["device"] = "cuda"
+        info["vram_gb"] = round(torch.cuda.get_device_properties(0).total_memory / 1e9, 1)
+        info["gpu_name"] = torch.cuda.get_device_name(0)
+        info["bf16"] = torch.cuda.is_bf16_supported()
+    return info
+
+HW = detect_hardware()
+
+def hw_banner() -> str:
+    if HW["device"] == "cuda":
+        tier = "🟢 GPU" if HW["vram_gb"] >= 16 else "🟡 GPU (pequena)"
+        return (
+            f"{tier} · {HW['gpu_name']} · {HW['vram_gb']} GB VRAM  |  "
+            f"BF16: {'✅' if HW['bf16'] else '❌'}  |  "
+            f"PEFT/LoRA: {'✅' if PEFT_AVAILABLE else '❌ (instale peft)'}"
+        )
+    return (
+        f"🔵 CPU  |  Threads: {torch.get_num_threads()}  |  "
+        f"PEFT/LoRA: {'✅' if PEFT_AVAILABLE else '❌'}"
+    )
+
+# ──────────────────────────────────────────────────────────────
+#  MODELO CATALOG  (modelo_id, max_vram_recomendado_gb)
+# ──────────────────────────────────────────────────────────────
+
+MODEL_CATALOG = {
+    # Tiny — roda até em CPU
+    "TinyLlama 1.1B": ("TinyLlama/TinyLlama-1.1B-Chat-v1.0", 4),
+    "SmolLM 1.7B": ("HuggingFaceTB/SmolLM2-1.7B-Instruct", 6),
+    # Médio — T4 OK com QLoRA
+    "Mistral 7B": ("mistralai/Mistral-7B-Instruct-v0.2", 14),
+    "Llama 3.1 8B": ("meta-llama/Meta-Llama-3.1-8B-Instruct", 16),
+    "Gemma 2 9B": ("google/gemma-2-9b-it", 18),
+    # Grande — A10G / A100
+    "Llama 3.1 70B": ("meta-llama/Meta-Llama-3.1-70B-Instruct", 80),
+    "Mixtral 8x7B": ("mistralai/Mixtral-8x7B-Instruct-v0.1", 48),
+}
+
+def available_models() -> list[str]:
+    """Filtra modelos que cabem no hardware atual."""
+    vram = HW["vram_gb"] if HW["device"] == "cuda" else 2
+    return [name for name, (_, req) in MODEL_CATALOG.items() if req <= max(vram * 1.2, 6)]
+
+# ──────────────────────────────────────────────────────────────
+#  ESTADO GLOBAL DO TREINAMENTO
+# ──────────────────────────────────────────────────────────────
+
+class TrainingState:
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.running       = False
+        self.cancelled     = False
+        self.logs: list    = []
+        self.progress: int = 0
+        self.total_steps   = 0
+        self.model         = None
+        self.tokenizer     = None
+        self.output_dir    = Path("./trained_model")
+        self.error: Optional[str] = None
+
+    def log(self, msg: str):
+        ts = time.strftime("%H:%M:%S")
+        self.logs.append(f"[{ts}] {msg}")
+        logger.info(msg)
+
+    def log_box(self) -> str:
+        return "\n".join(self.logs[-60:])  # últimas 60 linhas
+
+STATE = TrainingState()
+
+# ──────────────────────────────────────────────────────────────
+#  CALLBACK PARA PROGRESSO EM TEMPO REAL
+# ──────────────────────────────────────────────────────────────
+
+class ProgressCallback(TrainerCallback):
+    def on_train_begin(self, args, state, control, **kwargs):
+        STATE.total_steps = state.max_steps
+        STATE.log(f"▶ Treinamento iniciado — {state.max_steps} steps")
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs:
+            loss = logs.get("loss", "—")
+            lr   = logs.get("learning_rate", "—")
+            step = state.global_step
+            STATE.progress = step
+            STATE.log(f"Step {step}/{STATE.total_steps}  loss={loss}  lr={lr}")
+
+    def on_step_end(self, args, state, control, **kwargs):
+        if STATE.cancelled:
+            control.should_training_stop = True
+
+    def on_train_end(self, args, state, control, **kwargs):
+        STATE.log("✅ Treinamento concluído!")
+
+# ──────────────────────────────────────────────────────────────
+#  DATASET HELPERS
+# ──────────────────────────────────────────────────────────────
+
+def load_user_dataset(source: str, hf_dataset: str, uploaded_file) -> Dataset:
+    """Carrega dataset de múltiplas fontes."""
+    if source == "HuggingFace Hub" and hf_dataset.strip():
+        ds = load_dataset(hf_dataset.strip(), split="train")
+        return ds
+
+    if source == "Upload CSV/JSONL" and uploaded_file is not None:
+        path = uploaded_file.name
+        if path.endswith(".csv"):
+            df = pd.read_csv(path)
+        else:
+            df = pd.read_json(path, lines=True)
+        return Dataset.from_pandas(df)
+
+    # Fallback: dataset de exemplo embutido
+    examples = [
+        {"text": "Instrução: Explique o que é machine learning.\nResposta: Machine learning é..."},
+        {"text": "Instrução: O que é uma rede neural?\nResposta: Uma rede neural é..."},
+        {"text": "Instrução: Como funciona o backpropagation?\nResposta: O backpropagation..."},
+    ]
+    return Dataset.from_list(examples)
+
+def tokenize_dataset(dataset: Dataset, tokenizer, max_length: int) -> Dataset:
+    text_col = next(
+        (c for c in ["text", "prompt", "instruction", "content"] if c in dataset.column_names),
+        dataset.column_names[0],
+    )
+
+    def tokenize(examples):
+        return tokenizer(
+            examples[text_col],
+            truncation=True,
+            max_length=max_length,
+            padding="max_length",
+        )
+
+    return dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)
+
+# ──────────────────────────────────────────────────────────────
+#  CORE: CARREGA MODELO
+# ──────────────────────────────────────────────────────────────
+
+def load_model_and_tokenizer(model_name: str, ft_mode: str):
+    model_id, _ = MODEL_CATALOG[model_name]
+    STATE.log(f"⬇ Carregando tokenizer: {model_id}")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Configuração de quantização
+    bnb_cfg = None
+    load_in_4bit = False
+    load_in_8bit = False
+
+    if HW["device"] == "cuda" and ft_mode == "QLoRA":
+        if not PEFT_AVAILABLE:
+            raise RuntimeError("Instale `peft` e `bitsandbytes` para QLoRA.")
+        STATE.log("🔧 Configurando 4-bit NF4 (QLoRA)...")
+        bnb_cfg = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16 if HW["bf16"] else torch.float16,
+        )
+        load_in_4bit = True
+
+    STATE.log(f"⬇ Carregando modelo ({ft_mode})...")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        quantization_config=bnb_cfg,
+        device_map="auto" if HW["device"] == "cuda" else None,
+        torch_dtype=torch.float16 if (HW["device"] == "cuda" and not HW["bf16"]) else "auto",
+        trust_remote_code=True,
+    )
+
+    if ft_mode in ("LoRA", "QLoRA") and PEFT_AVAILABLE:
+        if load_in_4bit:
+            model = prepare_model_for_kbit_training(model)
+
+        lora_cfg = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            lora_dropout=0.05,
+            bias="none",
+            task_type=TaskType.CAUSAL_LM,
+            target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
+        )
+        model = get_peft_model(model, lora_cfg)
+        model.print_trainable_parameters()
+        trainable, total, pct = model.get_nb_trainable_parameters()
+        STATE.log(f"📊 Parâmetros treináveis: {trainable:,} / {total:,} ({pct:.2f}%)")
+
+    elif ft_mode == "Full Fine-Tuning":
+        STATE.log("⚠ Full fine-tuning: todos os pesos serão atualizados.")
+        if HW["device"] != "cuda":
+            STATE.log("⚠ Full fine-tuning em CPU será MUITO lento.")
+
+    return model, tokenizer
+
+# ──────────────────────────────────────────────────────────────
+#  CORE: TREINAMENTO
+# ──────────────────────────────────────────────────────────────
+
+def run_training(
+    model_name, ft_mode,
+    dataset_source, hf_dataset, uploaded_file,
+    epochs, batch_size, learning_rate, max_length,
+    warmup_steps, weight_decay, grad_accum,
+):
+    try:
+        STATE.reset()
+        STATE.running = True
+        STATE.output_dir = Path(f"./trained_{model_name.replace(' ', '_')}_{ft_mode}")
+        STATE.output_dir.mkdir(parents=True, exist_ok=True)
+
+        STATE.log(f"🖥 Hardware: {HW['gpu_name'] if HW['device']=='cuda' else 'CPU'}")
+        STATE.log(f"📦 Modelo: {model_name}  |  Modo: {ft_mode}")
+
+        # 1. Dataset
+        STATE.log("📂 Carregando dataset...")
+        raw_ds = load_user_dataset(dataset_source, hf_dataset, uploaded_file)
+        STATE.log(f"✅ Dataset: {len(raw_ds)} exemplos")
+
+        # 2. Modelo
+        model, tokenizer = load_model_and_tokenizer(model_name, ft_mode)
+        STATE.tokenizer = tokenizer
+
+        # 3. Tokenização
+        STATE.log("🔤 Tokenizando dataset...")
+        tokenized = tokenize_dataset(raw_ds, tokenizer, max_length)
+        tokenized = tokenized.train_test_split(test_size=0.05, seed=42)
+
+        # 4. TrainingArguments
+        use_fp16 = HW["device"] == "cuda" and not HW["bf16"]
+        use_bf16 = HW["device"] == "cuda" and HW["bf16"]
+
+        args = TrainingArguments(
+            output_dir=str(STATE.output_dir),
+            num_train_epochs=epochs,
+            per_device_train_batch_size=batch_size,
+            gradient_accumulation_steps=grad_accum,
+            learning_rate=learning_rate,
+            warmup_steps=warmup_steps,
+            weight_decay=weight_decay,
+            fp16=use_fp16,
+            bf16=use_bf16,
+            logging_steps=5,
+            save_steps=50,
+            save_total_limit=2,
+            eval_strategy="steps",
+            eval_steps=50,
+            load_best_model_at_end=True,
+            report_to="none",
+            dataloader_pin_memory=(HW["device"] == "cuda"),
+        )
+
+        collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+        trainer = Trainer(
+            model=model,
+            args=args,
+            train_dataset=tokenized["train"],
+            eval_dataset=tokenized["test"],
+            data_collator=collator,
+            callbacks=[ProgressCallback()],
+        )
+
+        STATE.log("🚀 Iniciando treinamento...")
+        trainer.train()
+
+        if not STATE.cancelled:
+            STATE.log("💾 Salvando modelo...")
+            model.save_pretrained(str(STATE.output_dir))
+            tokenizer.save_pretrained(str(STATE.output_dir))
+
+            # Salva metadados
+            meta = {
+                "base_model": MODEL_CATALOG[model_name][0],
+                "ft_mode": ft_mode,
+                "epochs": epochs,
+                "learning_rate": learning_rate,
+                "dataset_source": dataset_source,
+                "hardware": HW,
+            }
+            (STATE.output_dir / "training_meta.json").write_text(json.dumps(meta, indent=2))
+
+            STATE.model = model
+            STATE.log(f"🎉 Modelo salvo em: {STATE.output_dir}")
+
+    except Exception as e:
+        STATE.error = str(e)
+        STATE.log(f"❌ Erro: {e}")
+        logger.exception(e)
+    finally:
+        STATE.running = False
+
+# ──────────────────────────────────────────────────────────────
+#  CHAT COM MODELO TREINADO
+# ──────────────────────────────────────────────────────────────
+
+def chat_with_model(message: str, history: list, max_new_tokens: int, temperature: float):
+    if STATE.model is None or STATE.tokenizer is None:
+        return history + [[message, "⚠ Nenhum modelo treinado disponível. Complete o treinamento primeiro."]]
+
+    prompt = ""
+    for user_msg, bot_msg in history:
+        prompt += f"Usuário: {user_msg}\nAssistente: {bot_msg}\n"
+    prompt += f"Usuário: {message}\nAssistente:"
+
+    inputs = STATE.tokenizer(prompt, return_tensors="pt")
+    if HW["device"] == "cuda":
+        inputs = {k: v.cuda() for k, v in inputs.items()}
+
+    with torch.no_grad():
+        outputs = STATE.model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            do_sample=temperature > 0,
+            pad_token_id=STATE.tokenizer.eos_token_id,
+        )
+
+    gen_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+    response = STATE.tokenizer.decode(gen_tokens, skip_special_tokens=True).strip()
+    return history + [[message, response]]
+
+# ──────────────────────────────────────────────────────────────
+#  ZIP & DOWNLOAD
+# ──────────────────────────────────────────────────────────────
+
+def create_download_zip() -> Optional[str]:
+    if not STATE.output_dir.exists():
+        return None
+    zip_path = Path("./model_export.zip")
+    shutil.make_archive("model_export", "zip", str(STATE.output_dir))
+    return str(zip_path)
+
+# ──────────────────────────────────────────────────────────────
+#  GRADIO UI
+# ──────────────────────────────────────────────────────────────
+
+CSS = """
+@import url('https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;700&family=Syne:wght@400;600;800&display=swap');
+
+* { box-sizing: border-box; }
+
+body, .gradio-container {
+    background: #0a0a0f !important;
+    color: #e8e6f0 !important;
+    font-family: 'Syne', sans-serif !important;
+}
+
+.gr-panel, .gr-box, .gr-block { background: transparent !important; }
+
+/* Header */
+.studio-header {
+    text-align: center;
+    padding: 2rem 0 1rem;
+    border-bottom: 1px solid #2a2a3a;
+    margin-bottom: 1.5rem;
+}
+.studio-header h1 {
+    font-family: 'Syne', sans-serif;
+    font-weight: 800;
+    font-size: 2.2rem;
+    letter-spacing: -0.02em;
+    color: #fff;
+    margin: 0;
+}
+.studio-header h1 span { color: #7c6af7; }
+.studio-header p {
+    font-family: 'Space Mono', monospace;
+    font-size: 0.75rem;
+    color: #6b6888;
+    margin-top: 0.4rem;
+    letter-spacing: 0.08em;
+}
+
+/* Hardware badge */
+.hw-badge {
+    font-family: 'Space Mono', monospace;
+    font-size: 0.72rem;
+    background: #12121e;
+    border: 1px solid #2a2a3a;
+    border-radius: 6px;
+    padding: 0.5rem 1rem;
+    color: #8a88a8;
+    text-align: center;
+    margin-bottom: 1.2rem;
+}
+
+/* Tabs */
+.tab-nav button {
+    font-family: 'Syne', sans-serif !important;
+    font-weight: 600 !important;
+    font-size: 0.85rem !important;
+    letter-spacing: 0.04em !important;
+    color: #6b6888 !important;
+    background: transparent !important;
+    border: none !important;
+    border-bottom: 2px solid transparent !important;
+    padding: 0.5rem 1.2rem !important;
+}
+.tab-nav button.selected {
+    color: #7c6af7 !important;
+    border-bottom-color: #7c6af7 !important;
+}
+
+/* Inputs */
+.gr-input, .gr-dropdown select, textarea {
+    background: #12121e !important;
+    border: 1px solid #2a2a3a !important;
+    color: #e8e6f0 !important;
+    border-radius: 8px !important;
+    font-family: 'Space Mono', monospace !important;
+    font-size: 0.8rem !important;
+}
+.gr-input:focus, textarea:focus {
+    border-color: #7c6af7 !important;
+    box-shadow: 0 0 0 2px rgba(124,106,247,0.15) !important;
+}
+
+/* Buttons */
+.gr-button {
+    font-family: 'Syne', sans-serif !important;
+    font-weight: 600 !important;
+    border-radius: 8px !important;
+    transition: all 0.15s !important;
+}
+.gr-button.primary {
+    background: #7c6af7 !important;
+    border: none !important;
+    color: #fff !important;
+}
+.gr-button.primary:hover { background: #6a58e0 !important; transform: translateY(-1px); }
+.gr-button.secondary {
+    background: transparent !important;
+    border: 1px solid #2a2a3a !important;
+    color: #8a88a8 !important;
+}
+.gr-button.stop { background: #c0392b !important; color: #fff !important; border: none !important; }
+
+/* Log box */
+.log-box textarea {
+    font-family: 'Space Mono', monospace !important;
+    font-size: 0.72rem !important;
+    line-height: 1.6 !important;
+    background: #07070f !important;
+    border: 1px solid #1e1e2e !important;
+    color: #a8e6cf !important;
+}
+
+/* Progress bar */
+.progress-bar-wrap .progress-bar { background: #7c6af7 !important; }
+
+/* Slider labels */
+.gr-form label {
+    font-family: 'Syne', sans-serif !important;
+    font-size: 0.82rem !important;
+    color: #8a88a8 !important;
+    font-weight: 600 !important;
+}
+
+/* Section labels */
+.section-label {
+    font-family: 'Space Mono', monospace;
+    font-size: 0.65rem;
+    letter-spacing: 0.12em;
+    color: #4a4868;
+    text-transform: uppercase;
+    margin: 1rem 0 0.4rem;
+}
+
+/* Chat bubbles */
+.message.user div { background: #1e1e30 !important; border-radius: 10px !important; }
+.message.bot div  { background: #12121e !important; border-radius: 10px !important; border: 1px solid #2a2a3a !important; }
+"""
+
+def build_ui():
+    models = available_models()
+    if not models:
+        models = list(MODEL_CATALOG.keys())[:2]  # fallback
+
+    with gr.Blocks(css=CSS, title="Fine-Tuning Studio") as demo:
+
+        # ── HEADER ────────────────────────────────────────────
+        gr.HTML(f"""
+        <div class="studio-header">
+            <h1>🧪 Fine-Tuning <span>Studio</span></h1>
+            <p>TREINE · CONVERTA · CONVERSE · EXPORTE</p>
+        </div>
+        <div class="hw-badge">{hw_banner()}</div>
+        """)
+
+        # ── TABS ──────────────────────────────────────────────
+        with gr.Tabs(elem_classes="tab-nav"):
+
+            # ════════════════════════════════
+            #  TAB 1 — CONFIGURAR & TREINAR
+            # ════════════════════════════════
+            with gr.Tab("⚙️ Treinar"):
+                with gr.Row():
+
+                    # Coluna esquerda — config
+                    with gr.Column(scale=1):
+                        gr.HTML('<div class="section-label">modelo</div>')
+                        model_dd = gr.Dropdown(
+                            choices=models,
+                            value=models[0],
+                            label="Modelo base",
+                            interactive=True,
+                        )
+                        ft_mode_dd = gr.Dropdown(
+                            choices=["LoRA", "QLoRA", "Full Fine-Tuning"],
+                            value="LoRA" if PEFT_AVAILABLE else "Full Fine-Tuning",
+                            label="Modo de fine-tuning",
+                        )
+
+                        gr.HTML('<div class="section-label">dataset</div>')
+                        ds_source = gr.Radio(
+                            choices=["HuggingFace Hub", "Upload CSV/JSONL", "Exemplo embutido"],
+                            value="Exemplo embutido",
+                            label="Fonte do dataset",
+                        )
+                        hf_ds_input = gr.Textbox(
+                            placeholder="ex: tatsu-lab/alpaca",
+                            label="Dataset ID (Hub)",
+                            visible=False,
+                        )
+                        upload_file = gr.File(
+                            label="CSV ou JSONL",
+                            file_types=[".csv", ".jsonl"],
+                            visible=False,
+                        )
+
+                        def toggle_ds(source):
+                            return (
+                                gr.update(visible=source == "HuggingFace Hub"),
+                                gr.update(visible=source == "Upload CSV/JSONL"),
+                            )
+                        ds_source.change(toggle_ds, ds_source, [hf_ds_input, upload_file])
+
+                        gr.HTML('<div class="section-label">hiperparâmetros</div>')
+                        epochs_sl    = gr.Slider(1, 10, value=3, step=1, label="Épocas")
+                        batch_sl     = gr.Slider(1, 16, value=2, step=1, label="Batch size")
+                        lr_sl        = gr.Slider(1e-5, 5e-4, value=2e-4, step=1e-5, label="Learning rate")
+                        max_len_sl   = gr.Slider(64, 2048, value=512, step=64, label="Max length (tokens)")
+                        grad_acc_sl  = gr.Slider(1, 16, value=4, step=1, label="Grad. accumulation")
+                        warmup_sl    = gr.Slider(0, 200, value=10, step=5, label="Warmup steps")
+                        wd_sl        = gr.Slider(0, 0.1, value=0.01, step=0.005, label="Weight decay")
+
+                    # Coluna direita — logs
+                    with gr.Column(scale=1):
+                        gr.HTML('<div class="section-label">log de treinamento</div>')
+                        log_box = gr.Textbox(
+                            label="",
+                            lines=24,
+                            max_lines=24,
+                            interactive=False,
+                            elem_classes="log-box",
+                            placeholder="O log aparecerá aqui quando o treinamento iniciar...",
+                        )
+                        progress = gr.Slider(
+                            0, 100, value=0, label="Progresso (%)", interactive=False
+                        )
+
+                with gr.Row():
+                    train_btn  = gr.Button("🚀 Iniciar Treinamento", variant="primary", scale=3)
+                    cancel_btn = gr.Button("⏹ Cancelar", variant="stop", scale=1)
+
+                status_md = gr.Markdown("")
+
+                # ── Lógica de treinamento ──────────────────
+                def start_training(
+                    model_name, ft_mode,
+                    ds_source, hf_ds, up_file,
+                    epochs, batch, lr, max_len,
+                    warmup, wd, grad_acc,
+                ):
+                    if STATE.running:
+                        return "⚠ Treinamento já em andamento."
+
+                    thread = threading.Thread(
+                        target=run_training,
+                        args=(
+                            model_name, ft_mode,
+                            ds_source, hf_ds, up_file,
+                            epochs, batch, lr, max_len,
+                            warmup, wd, grad_acc,
+                        ),
+                        daemon=True,
+                    )
+                    thread.start()
+                    return "▶ Treinamento iniciado..."
+
+                def cancel_training():
+                    STATE.cancelled = True
+                    return "⏹ Cancelamento solicitado."
+
+                def poll_logs():
+                    """Polling a cada 2s para atualizar log e progresso."""
+                    while True:
+                        pct = 0
+                        if STATE.total_steps > 0:
+                            pct = min(100, int(STATE.progress / STATE.total_steps * 100))
+                        yield STATE.log_box(), pct
+                        time.sleep(2)
+
+                train_btn.click(
+                    start_training,
+                    inputs=[
+                        model_dd, ft_mode_dd,
+                        ds_source, hf_ds_input, upload_file,
+                        epochs_sl, batch_sl, lr_sl, max_len_sl,
+                        warmup_sl, wd_sl, grad_acc_sl,
+                    ],
+                    outputs=status_md,
+                )
+                cancel_btn.click(cancel_training, outputs=status_md)
+
+                demo.load(poll_logs, outputs=[log_box, progress], every=2)
+
+            # ════════════════════════════════
+            #  TAB 2 — CHAT
+            # ════════════════════════════════
+            with gr.Tab("💬 Chat"):
+                gr.Markdown(
+                    "**Converse com o modelo treinado.** Complete o treinamento na aba anterior primeiro.",
+                    elem_id="chat-hint",
+                )
+                chatbot = gr.Chatbot(height=440, label="Conversa")
+
+                with gr.Row():
+                    chat_input = gr.Textbox(
+                        placeholder="Digite sua mensagem...",
+                        label="",
+                        scale=4,
+                    )
+                    send_btn = gr.Button("Enviar", variant="primary", scale=1)
+
+                with gr.Accordion("⚙️ Parâmetros de geração", open=False):
+                    max_new_sl = gr.Slider(32, 1024, value=256, step=32, label="Max new tokens")
+                    temp_sl    = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature")
+
+                clear_btn = gr.Button("🗑 Limpar conversa", variant="secondary")
+
+                send_btn.click(
+                    chat_with_model,
+                    inputs=[chat_input, chatbot, max_new_sl, temp_sl],
+                    outputs=chatbot,
+                )
+                chat_input.submit(
+                    chat_with_model,
+                    inputs=[chat_input, chatbot, max_new_sl, temp_sl],
+                    outputs=chatbot,
+                )
+                clear_btn.click(lambda: [], outputs=chatbot)
+
+            # ════════════════════════════════
+            #  TAB 3 — EXPORTAR
+            # ════════════════════════════════
+            with gr.Tab("📦 Exportar"):
+                gr.Markdown("### Download dos pesos treinados")
+                gr.Markdown(
+                    "Após o treinamento, clique abaixo para gerar um `.zip` com todos os pesos e metadados."
+                )
+
+                with gr.Row():
+                    zip_btn      = gr.Button("📦 Gerar ZIP", variant="primary")
+                    download_out = gr.File(label="Download", interactive=False)
+
+                export_status = gr.Markdown("")
+
+                def generate_zip():
+                    path = create_download_zip()
+                    if path:
+                        return path, "✅ ZIP gerado! Clique para baixar."
+                    return None, "⚠ Nenhum modelo treinado encontrado. Complete o treinamento primeiro."
+
+                zip_btn.click(generate_zip, outputs=[download_out, export_status])
+
+                gr.Markdown("---")
+                gr.Markdown("### Push para HuggingFace Hub")
+                gr.Markdown(
+                    "Para fazer push do modelo para o Hub, configure o `HF_TOKEN` nas **Secrets** do Space "
+                    "e use `model.push_to_hub('seu-usuario/nome-do-modelo')` no terminal."
+                )
+
+                # Metadados do treino
+                gr.HTML('<div class="section-label">metadados do treino</div>')
+
+                def get_meta():
+                    meta_file = STATE.output_dir / "training_meta.json"
+                    if meta_file.exists():
+                        return meta_file.read_text()
+                    return "Sem metadados ainda."
+
+                meta_box = gr.Code(label="training_meta.json", language="json", interactive=False)
+                refresh_meta_btn = gr.Button("🔄 Atualizar metadados", variant="secondary")
+                refresh_meta_btn.click(get_meta, outputs=meta_box)
+
+        # Footer
+        gr.HTML("""
+        <div style="text-align:center; margin-top:2rem; font-family:'Space Mono',monospace;
+                    font-size:0.65rem; color:#3a3858; letter-spacing:0.1em;">
+            FINE-TUNING STUDIO · HUGGINGFACE SPACE · ADAPTA-SE AO HARDWARE DISPONÍVEL
+        </div>
+        """)
+
+    return demo
+
+
+# ──────────────────────────────────────────────────────────────
+#  ENTRY POINT
+# ──────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    app = build_ui()
+    app.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True,
+    )