Spaces:
Runtime error
Runtime error
File size: 3,176 Bytes
b7ac716 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | import os
import subprocess
import sys
import threading
from pathlib import Path
import gradio as gr
LOG_FILE = Path("/tmp/apochat_training.log")
STATUS_FILE = Path("/tmp/apochat_training_status.txt")
DATASET = os.environ.get("DATASET", "apoapps/apochat-gemma4-e2b-chat-v1")
OUTPUT_REPO = os.environ.get("OUTPUT_REPO", "apoapps/apochat-gemma4-e2b-apochat-tuned-v2")
def _status() -> str:
if STATUS_FILE.exists():
return STATUS_FILE.read_text().strip()
return "idle"
def _log_tail(n: int = 80) -> str:
if not LOG_FILE.exists():
return "No logs yet."
lines = LOG_FILE.read_text().splitlines()
return "\n".join(lines[-n:])
def run_training(use_qlora: bool, epochs: float, learning_rate: float) -> None:
STATUS_FILE.write_text("running")
LOG_FILE.write_text("")
cmd = [
sys.executable,
"finetune_apochat_peft.py",
"--dataset", DATASET,
"--output-dir", "/tmp/apochat-peft-output",
"--push-to-hub", OUTPUT_REPO,
"--epochs", str(epochs),
"--learning-rate", str(learning_rate),
]
if use_qlora:
cmd.append("--use-qlora")
with open(LOG_FILE, "a") as log_f:
log_f.write(f"Running: {' '.join(cmd)}\n")
proc = subprocess.Popen(
cmd,
stdout=log_f,
stderr=subprocess.STDOUT,
cwd=str(Path(__file__).parent),
)
proc.wait()
STATUS_FILE.write_text("done" if proc.returncode == 0 else f"failed:{proc.returncode}")
def start_training(use_qlora: bool, epochs: float, learning_rate: float):
if _status() == "running":
return "already running", _log_tail()
threading.Thread(target=run_training, args=(use_qlora, epochs, learning_rate), daemon=True).start()
return "started", _log_tail()
def refresh() -> tuple[str, str]:
return _status(), _log_tail()
with gr.Blocks(title="Apochat Gemma 4 E2B Trainer") as demo:
gr.Markdown("""
# Apochat Gemma 4 E2B Trainer
Fine-tune the base `google/gemma-4-E2B-it` model on the Apochat chat dataset using
PEFT/LoRA (optionally QLoRA). Training is started manually after you upgrade the
Space hardware to a GPU.
**Before clicking Start:** go to the Space settings and set Hardware to a GPU
(e.g. `t4-small`, `a10g-small`, `l4x1`). Free `cpu-basic` will not finish training.
""")
with gr.Row():
use_qlora = gr.Checkbox(label="Use QLoRA (4-bit, saves VRAM)", value=True)
epochs = gr.Number(label="Epochs", value=1.0, minimum=0.1, maximum=5.0)
learning_rate = gr.Number(label="Learning rate", value=2e-4, minimum=1e-5, maximum=1e-3)
start_btn = gr.Button("Start training", variant="primary")
status_box = gr.Textbox(label="Status", value=_status())
log_box = gr.Textbox(label="Log", lines=25, value=_log_tail(), max_lines=25)
start_btn.click(start_training, inputs=[use_qlora, epochs, learning_rate], outputs=[status_box, log_box])
refresh_btn = gr.Button("Refresh")
refresh_btn.click(refresh, outputs=[status_box, log_box])
demo.load(refresh, outputs=[status_box, log_box], every=10)
if __name__ == "__main__":
demo.launch()
|