Update app.py

app.py

@@ -1,25 +1,104 @@
+import os
+import time
+from pathlib import Path
+
 import gradio as gr
-from trainer import run_finetune
 
+from src.train import finetune_lora
+from src.infer import load_generator, generate_text
+
+def _default_output_root() -> Path:
+    # On Spaces, /data exists if Persistent Storage is enabled.
+    # Otherwise fall back to repo-local outputs/.
+    return Path("/data/outputs") if Path("/data").exists() else Path("outputs")
+
+def run_train(
+    base_model: str,
+    dataset_id: str,
+    text_column: str,
+    max_train_samples: int,
+    max_steps: int,
+    lr: float,
+    batch_size: int,
+    lora_r: int,
+    lora_alpha: int,
+    lora_dropout: float,
+):
+    out_root = _default_output_root()
+    run_id = time.strftime("%Y%m%d-%H%M%S")
+    out_dir = out_root / run_id
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    status = finetune_lora(
+        base_model=base_model.strip(),
+        dataset_id=dataset_id.strip(),
+        text_column=text_column.strip(),
+        output_dir=str(out_dir),
+        max_train_samples=max_train_samples,
+        max_steps=max_steps,
+        learning_rate=lr,
+        batch_size=batch_size,
+        lora_r=lora_r,
+        lora_alpha=lora_alpha,
+        lora_dropout=lora_dropout,
+    )
 
-def start_training(base_model, dataset_file, epochs):
-    output = run_finetune(
-        base_model=base_model,
-        dataset_path=dataset_file.name,
-        epochs=int(epochs),
+    adapter_path = out_dir / "adapter"
+    return (
+        f"Done.\n\nSaved to: {out_dir}\n\n{status}",
+        str(adapter_path) if adapter_path.exists() else None,
+        str(out_dir),
     )
-    return output
 
+def run_generate(base_model: str, adapter_dir: str, prompt: str, max_new_tokens: int, temperature: float):
+    gen = load_generator(base_model.strip(), adapter_dir.strip())
+    return generate_text(gen, prompt, max_new_tokens=max_new_tokens, temperature=temperature)
+
+with gr.Blocks(title="Fine-tune Pipeline (Docker)") as demo:
+    gr.Markdown("# Fine-tuning pipeline (LoRA) — Docker Space")
+
+    with gr.Tab("Train"):
+        base_model = gr.Textbox(value="sshleifer/tiny-gpt2", label="Base model (HF Hub id)")
+        dataset_id = gr.Textbox(value="karpathy/tiny_shakespeare", label="Dataset (HF Hub id)")
+        text_column = gr.Textbox(value="text", label="Text column")
+
+        with gr.Row():
+            max_train_samples = gr.Number(value=2000, precision=0, label="Max train samples")
+            max_steps = gr.Number(value=100, precision=0, label="Max steps")
+
+        with gr.Row():
+            lr = gr.Number(value=2e-4, label="Learning rate")
+            batch_size = gr.Number(value=2, precision=0, label="Batch size")
+
+        with gr.Row():
+            lora_r = gr.Number(value=8, precision=0, label="LoRA r")
+            lora_alpha = gr.Number(value=16, precision=0, label="LoRA alpha")
+            lora_dropout = gr.Number(value=0.05, label="LoRA dropout")
+
+        train_btn = gr.Button("Start fine-tune")
+        train_out = gr.Textbox(lines=10, label="Status")
+        adapter_file = gr.File(label="Adapter folder (download)")
+        out_dir_box = gr.Textbox(label="Output directory")
+
+        train_btn.click(
+            fn=run_train,
+            inputs=[base_model, dataset_id, text_column, max_train_samples, max_steps, lr, batch_size, lora_r, lora_alpha, lora_dropout],
+            outputs=[train_out, adapter_file, out_dir_box],
+            queue=True,
+        )
+
+    with gr.Tab("Generate"):
+        base_model2 = gr.Textbox(value="sshleifer/tiny-gpt2", label="Base model (must match training)")
+        adapter_dir = gr.Textbox(placeholder="Paste the output adapter dir path (e.g., outputs/2026.../adapter)", label="Adapter directory")
+        prompt = gr.Textbox(value="To be, or not to be,", lines=3, label="Prompt")
+
+        with gr.Row():
+            max_new_tokens = gr.Slider(16, 256, value=80, step=1, label="Max new tokens")
+            temperature = gr.Slider(0.1, 1.5, value=0.9, step=0.05, label="Temperature")
+
+        gen_btn = gr.Button("Generate")
+        gen_out = gr.Textbox(lines=10, label="Output")
 
-ui = gr.Interface(
-    fn=start_training,
-    inputs=[
-        gr.Textbox(value="distilbert-base-uncased", label="Base model"),
-        gr.File(label="Dataset (jsonl)"),
-        gr.Number(value=3, label="Epochs")
-    ],
-    outputs=gr.Textbox(),
-    title="HuggingFace Fine-Tuning Space"
-)
+        gen_btn.click(fn=run_generate, inputs=[base_model2, adapter_dir, prompt, max_new_tokens, temperature], outputs=[gen_out])
 
-ui.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch()