Create app.py

app.py

@@ -0,0 +1,279 @@
+import os
+import gradio as gr
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling,
+)
+from peft import LoraConfig, get_peft_model, TaskType
+import threading
+
+# ── Globals ──────────────────────────────────────────────────────────────────
+training_log = []
+training_thread = None
+stop_flag = threading.Event()
+
+
+def log(msg: str):
+    training_log.append(msg)
+    print(msg)
+
+
+# ── Core training function ────────────────────────────────────────────────────
+def run_finetuning(
+    model_name: str,
+    dataset_name: str,
+    dataset_config: str,
+    text_column: str,
+    num_train_epochs: int,
+    per_device_batch_size: int,
+    learning_rate: float,
+    max_seq_length: int,
+    use_lora: bool,
+    lora_r: int,
+    output_dir: str,
+):
+    global training_log, stop_flag
+    training_log = []
+    stop_flag.clear()
+
+    try:
+        log(f"🔧 Loading tokenizer: {model_name}")
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        log(f"📦 Loading model: {model_name}")
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,  # CPU-safe
+            low_cpu_mem_usage=True,
+        )
+
+        if use_lora:
+            log(f"⚡ Applying LoRA (r={lora_r}) ...")
+            lora_config = LoraConfig(
+                task_type=TaskType.CAUSAL_LM,
+                r=lora_r,
+                lora_alpha=lora_r * 2,
+                lora_dropout=0.05,
+                bias="none",
+                target_modules=["c_attn", "c_proj", "q_proj", "v_proj", "k_proj", "o_proj"],
+            )
+            model = get_peft_model(model, lora_config)
+            trainable, total = model.get_nb_trainable_parameters()
+            log(f"   Trainable params: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)")
+
+        log(f"📂 Loading dataset: {dataset_name}" + (f" ({dataset_config})" if dataset_config else ""))
+        ds_kwargs = {"split": "train", "trust_remote_code": True}
+        if dataset_config.strip():
+            dataset = load_dataset(dataset_name, dataset_config, **ds_kwargs)
+        else:
+            dataset = load_dataset(dataset_name, **ds_kwargs)
+
+        # Take a small sample for demo / CPU friendliness
+        dataset = dataset.select(range(min(500, len(dataset))))
+        log(f"   Using {len(dataset)} training samples")
+
+        def tokenize(batch):
+            texts = [str(t) for t in batch[text_column]]
+            return tokenizer(
+                texts,
+                truncation=True,
+                max_length=max_seq_length,
+                padding="max_length",
+            )
+
+        log("🔤 Tokenizing dataset ...")
+        tokenized = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)
+        tokenized.set_format("torch")
+
+        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=num_train_epochs,
+            per_device_train_batch_size=per_device_batch_size,
+            learning_rate=learning_rate,
+            logging_steps=5,
+            save_strategy="epoch",
+            fp16=False,
+            bf16=False,
+            no_cuda=True,
+            report_to="none",
+            disable_tqdm=False,
+        )
+
+        class LogCallback(torch.utils.data.Dataset):
+            pass
+
+        from transformers import TrainerCallback
+
+        class StreamLogger(TrainerCallback):
+            def on_log(self, args, state, control, logs=None, **kwargs):
+                if logs:
+                    step = state.global_step
+                    loss = logs.get("loss", "—")
+                    lr = logs.get("learning_rate", "—")
+                    log(f"  step {step:>4} | loss: {loss} | lr: {lr}")
+
+            def on_epoch_end(self, args, state, control, **kwargs):
+                log(f"✅ Epoch {int(state.epoch)} complete")
+                if stop_flag.is_set():
+                    control.should_training_stop = True
+
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized,
+            data_collator=data_collator,
+            callbacks=[StreamLogger()],
+        )
+
+        log("🚀 Starting training ...")
+        trainer.train()
+
+        log(f"💾 Saving model to: {output_dir}")
+        trainer.save_model(output_dir)
+        tokenizer.save_pretrained(output_dir)
+        log("🎉 Fine-tuning complete!")
+
+    except Exception as e:
+        log(f"❌ Error: {e}")
+
+
+# ── Gradio helpers ────────────────────────────────────────────────────────────
+def start_training(
+    model_name, dataset_name, dataset_config, text_column,
+    num_epochs, batch_size, learning_rate, max_seq_len,
+    use_lora, lora_r, output_dir,
+):
+    global training_thread
+    if training_thread and training_thread.is_alive():
+        return "⚠️ Training already running!"
+
+    training_thread = threading.Thread(
+        target=run_finetuning,
+        args=(
+            model_name, dataset_name, dataset_config, text_column,
+            num_epochs, batch_size, learning_rate, max_seq_len,
+            use_lora, lora_r, output_dir,
+        ),
+        daemon=True,
+    )
+    training_thread.start()
+    return "Training started! Check the log below."
+
+
+def stop_training():
+    stop_flag.set()
+    return "🛑 Stop signal sent."
+
+
+def get_logs():
+    return "\n".join(training_log) if training_log else "No logs yet..."
+
+
+def is_running():
+    return "🟢 Running" if (training_thread and training_thread.is_alive()) else "⚫ Idle"
+
+
+# ── Gradio UI ─────────────────────────────────────────────────────────────────
+with gr.Blocks(
+    title="LLM Fine-Tuner",
+    theme=gr.themes.Base(
+        primary_hue="emerald",
+        neutral_hue="zinc",
+        font=gr.themes.GoogleFont("JetBrains Mono"),
+    ),
+    css="""
+    .container { max-width: 900px; margin: auto; }
+    .gr-button-primary { background: #10b981 !important; }
+    footer { display: none !important; }
+    """,
+) as demo:
+    gr.Markdown(
+        """
+# 🤖 LLM Fine-Tuner
+Fine-tune small language models on Hugging Face datasets — CPU-friendly with LoRA support.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### 🧠 Model")
+            model_name = gr.Dropdown(
+                choices=[
+                    "distilgpt2",
+                    "gpt2",
+                    "facebook/opt-125m",
+                    "EleutherAI/pythia-70m",
+                    "EleutherAI/pythia-160m",
+                    "microsoft/phi-1_5",
+                ],
+                value="distilgpt2",
+                label="Base Model",
+                allow_custom_value=True,
+            )
+
+            gr.Markdown("### 📦 Dataset")
+            dataset_name = gr.Textbox(value="wikitext", label="Dataset Name (HF Hub)")
+            dataset_config = gr.Textbox(value="wikitext-2-raw-v1", label="Dataset Config (optional)")
+            text_column = gr.Textbox(value="text", label="Text Column")
+
+        with gr.Column(scale=1):
+            gr.Markdown("### ⚙️ Training")
+            num_epochs = gr.Slider(1, 10, value=1, step=1, label="Epochs")
+            batch_size = gr.Slider(1, 16, value=2, step=1, label="Batch Size")
+            learning_rate = gr.Number(value=2e-4, label="Learning Rate")
+            max_seq_len = gr.Slider(32, 512, value=128, step=32, label="Max Sequence Length")
+            output_dir = gr.Textbox(value="./finetuned-model", label="Output Directory")
+
+            gr.Markdown("### ⚡ LoRA (recommended for CPU)")
+            use_lora = gr.Checkbox(value=True, label="Use LoRA")
+            lora_r = gr.Slider(4, 64, value=8, step=4, label="LoRA Rank (r)")
+
+    with gr.Row():
+        start_btn = gr.Button("🚀 Start Fine-Tuning", variant="primary")
+        stop_btn = gr.Button("🛑 Stop", variant="secondary")
+        status_btn = gr.Button("🔄 Refresh Status")
+
+    status_box = gr.Textbox(label="Status", value="⚫ Idle", interactive=False)
+    log_box = gr.Textbox(
+        label="Training Log",
+        lines=20,
+        max_lines=30,
+        interactive=False,
+        placeholder="Logs will appear here once training starts...",
+    )
+
+    start_btn.click(
+        fn=start_training,
+        inputs=[
+            model_name, dataset_name, dataset_config, text_column,
+            num_epochs, batch_size, learning_rate, max_seq_len,
+            use_lora, lora_r, output_dir,
+        ],
+        outputs=status_box,
+    )
+
+    stop_btn.click(fn=stop_training, outputs=status_box)
+    status_btn.click(fn=lambda: (is_running(), get_logs()), outputs=[status_box, log_box])
+
+    gr.Markdown(
+        """
+---
+**Tips:**
+- `distilgpt2` (82M) is the best starting point on CPU.  
+- Enable **LoRA** to drastically reduce memory and training time.  
+- Keep **Max Sequence Length ≤ 128** and **Batch Size = 1–2** on free CPU tier.  
+- The dataset is capped at **500 samples** for CPU-friendly runs — edit the code to increase.
+        """
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=False)