Update app.py

app.py

@@ -1,160 +1,150 @@
+# app.py
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
-from datasets import load_dataset
-import torch
-import os
-
-#-------------------------------Functions----------------------------------------------#
-def load_and_preprocess_data(dataset_name, tokenizer):
+from backend import verify_hf_token, get_user_runs, get_run_logs, queue_training_run, start_training_if_free
+from utils import ARCH_ANALOGIES, get_auto_hyperparams
+
+# ------------------------------ STATE ------------------------------
+user_state = {"user_id": None, "hf_token": "", "current_run_id": None}
+
+# ------------------------------ PAGES ------------------------------
+
+def page_login(hf_token):
+    user_id, msg = verify_hf_token(hf_token)
+    if user_id:
+        user_state["user_id"] = user_id
+        user_state["hf_token"] = hf_token
+        return gr.update(visible=False), gr.update(visible=True), msg
+    else:
+        return gr.update(), gr.update(), msg
+
+def page_processes():
+    runs = get_user_runs(user_state["user_id"])
+    run_list = "\n".join([
+        f"🍳 Run #{r[0]} | {r[1].upper()} x{r[2]} layers | {r[3]} | {r[4]}"
+        for r in runs
+    ]) or "No runs yet. Start cooking!"
+    return run_list
+
+def load_run_logs(run_id_str):
     try:
-        dataset = load_dataset(dataset_name, split="train")
-    except Exception as e:
-        return None, f"Error loading dataset: {e}"
-
-    def tokenize_function(examples):
-        return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
+        run_id = int(run_id_str)
+        logs, status = get_run_logs(run_id)
+        return f"Status: {status}\n\n{logs}"
+    except:
+        return "Invalid run ID."
+
+def page_architecture_next(arch_type, num_layers):
+    analogy = ARCH_ANALOGIES.get(arch_type, "")
+    auto_config = get_auto_hyperparams(arch_type, num_layers)
+    user_state["arch_config"] = {
+        "arch_type": arch_type,
+        "num_layers": num_layers,
+        "auto_config": auto_config
+    }
+    return (
+        gr.update(visible=False),
+        gr.update(visible=True),
+        f"🧠 {analogy}\n\nAuto-Seasoning™ Suggestion:\nLR: {auto_config['learning_rate']} | Epochs: {auto_config['epochs']} | Batch: {auto_config['batch_size']}"
+    )
 
-    try:
-        tokenized_datasets = dataset.map(tokenize_function, batched=True, num_proc=4, remove_columns=["text"])
-    except Exception as e:
-        return None, f"Error tokenizing dataset: {e}"
-
-    return tokenized_datasets, None
-#---------------------------------------------------------------------------------------#
-
-def train_model(architecture_size, api_key, repo_name, push_to_hub):
-    # Map architecture size to model name
-    model_name_mapping = {
-        "Small": "distilgpt2",
-        "Medium": "gpt2",
-        "Large": "gpt2-medium",
+def page_hyperparams_next(lr, epochs, batch_size):
+    config = user_state["arch_config"]
+    final_config = {
+        "arch_type": config["arch_type"],
+        "num_layers": config["num_layers"],
+        "learning_rate": float(lr) if lr else config["auto_config"]["learning_rate"],
+        "epochs": int(epochs) if epochs else config["auto_config"]["epochs"],
+        "batch_size": int(batch_size) if batch_size else config["auto_config"]["batch_size"],
     }
-    model_name = model_name_mapping[architecture_size]
+    run_id = queue_training_run(user_state["user_id"], final_config)
+    user_state["current_run_id"] = run_id
+    # Try to start if RAM allows
+    can_start = start_training_if_free()
+    status = "queued. Waiting for available stove 🔥..." if not can_start else "starting..."
+    return (
+        gr.update(visible=False),
+        gr.update(visible=True),
+        f"✅ Run #{run_id} {status}\nCheck 'Your Processes' for logs!"
+    )
 
-    # Device setup
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    device_msg = "CUDA is available! Training will be faster on GPU." if torch.cuda.is_available() else "CUDA not available. Training on CPU will be slow."
+# ------------------------------ UI ------------------------------
+
+with gr.Blocks(title="LLM Kitchen 🍳") as demo:
+    gr.Markdown("# 🍳 Welcome to LLM Kitchen")
+    gr.Markdown("### Cook your own language model — from scratch!")
+
+    # ---- PAGE 1: LOGIN ----
+    with gr.Group() as page_login_ui:
+        gr.Markdown("### 🔐 Step 1: Login with Hugging Face Token")
+        token_input = gr.Textbox(label="HF Token (starts with 'hf_')", type="password")
+        login_btn = gr.Button("Login to Kitchen", variant="primary")
+        login_msg = gr.Markdown()
+
+    # ---- PAGE 2: PROCESSES ----
+    with gr.Group(visible=False) as page_processes_ui:
+        gr.Markdown("### 🧑‍🍳 Your Processes")
+        refresh_btn = gr.Button("Refresh List")
+        runs_display = gr.Textbox(label="Your Training Runs", lines=8)
+        run_id_input = gr.Textbox(label="Enter Run ID to View Logs")
+        view_logs_btn = gr.Button("View Logs")
+        logs_display = gr.Textbox(label="Training Logs", lines=10)
+        new_run_btn = gr.Button("➕ Start New Process", variant="primary")
+
+    # ---- PAGE 3: ARCHITECTURE ----
+    with gr.Group(visible=False) as page_arch_ui:
+        gr.Markdown("### 🏗️ Step 2: Choose Your Architecture")
+        arch_dropdown = gr.Dropdown(["cnn", "rnn", "transformer"], label="Architecture Type")
+        layers_slider = gr.Slider(1, 16, value=4, step=1, label="Number of Layers (Think: # of sauce reductions)")
+        arch_next_btn = gr.Button("Next → Hyperparameters")
+        arch_analogy = gr.Markdown()
+
+    # ---- PAGE 4: HYPERPARAMETERS ----
+    with gr.Group(visible=False) as page_hyper_ui:
+        gr.Markdown("### 🧂 Step 3: Season Your Model (Hyperparameters)")
+        gr.Markdown("Use Auto-Seasoning™ or customize manually")
+        lr_input = gr.Number(label="Learning Rate (Saltiness)")
+        epochs_input = gr.Number(label="Epochs (Simmer Time)", precision=0)
+        batch_input = gr.Number(label="Batch Size (Spoon Size)", precision=0)
+        hyper_next_btn = gr.Button("Start Cooking! 🍲")
+
+    # ---- PAGE 5: TRAINING STARTED ----
+    with gr.Group(visible=False) as page_train_ui:
+        train_status = gr.Markdown("Starting your training run...")
+
+    # ------------------------------ EVENTS ------------------------------
+
+    login_btn.click(
+        page_login,
+        inputs=token_input,
+        outputs=[page_login_ui, page_processes_ui, login_msg]
+    )
 
-    # Validate push_to_hub inputs
-    if push_to_hub:
-        if not api_key or not api_key.strip():
-            return "❌ Error: You must provide a Hugging Face API key if pushing to hub is selected."
-        if not repo_name or not repo_name.strip():
-            return "❌ Error: You must provide a repository name if pushing to hub is selected."
+    refresh_btn.click(
+        page_processes,
+        outputs=runs_display
+    )
 
-    try:
-        # Load dataset
-        dataset_name = "wikitext-2-raw-v1"
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-
-        tokenized_datasets, error_msg = load_and_preprocess_data(dataset_name, tokenizer)
-        if error_msg:
-            return f"❌ {error_msg}"
-        if tokenized_datasets is None:
-            return "❌ Failed to load and preprocess dataset."
-
-        # Load model
-        model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
-        model.resize_token_embeddings(len(tokenizer))
-
-        # Training args
-        output_dir = "./results"
-        training_args = TrainingArguments(
-            output_dir=output_dir,
-            num_train_epochs=1,
-            per_device_train_batch_size=4,
-            save_steps=500,
-            save_total_limit=1,
-            logging_steps=250,
-            learning_rate=5e-5,
-            weight_decay=0.01,
-            push_to_hub=push_to_hub,
-            hub_model_id=repo_name if push_to_hub else None,
-            hub_token=api_key if push_to_hub else None,
-            fp16=torch.cuda.is_available(),
-        )
-
-        # Data collator
-        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
-
-        # Trainer
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=tokenized_datasets,
-            data_collator=data_collator,
-        )
-
-        # Train
-        trainer.train()
-
-        # Save locally
-        trainer.save_model(output_dir)
-
-        # Evaluate
-        eval_results = trainer.evaluate()
-        eval_loss = eval_results.get('eval_loss', 'N/A')
-
-        # Push to hub if selected
-        if push_to_hub:
-            trainer.push_to_hub()
-            hub_msg = f"✅ Model pushed to Hugging Face Hub: {repo_name}"
-        else:
-            hub_msg = "ℹ️ Model saved locally at ./results (not pushed to hub)."
-
-        return f"""✅ Training Complete!
-- Device: {device_msg}
-- Eval Loss: {eval_loss}
-- {hub_msg}
-"""
-
-    except Exception as e:
-        return f"❌ Training Error: {str(e)}"
-
-# ----------------------------- Gradio Interface ----------------------------- #
-
-with gr.Blocks(title="LLM Builder - Gradio") as demo:
-    gr.Markdown("# 🤖 LLM Builder")
-    gr.Markdown("### 1. Select Model Architecture")
-
-    architecture_size = gr.Dropdown(
-        choices=["Small", "Medium", "Large"],
-        value="Small",
-        label="Choose Model Size",
-        info="Select the size of the model. Larger models have more parameters."
+    view_logs_btn.click(
+        load_run_logs,
+        inputs=run_id_input,
+        outputs=logs_display
+    )
+
+    new_run_btn.click(
+        lambda: (gr.update(visible=False), gr.update(visible=True), ""),
+        outputs=[page_processes_ui, page_arch_ui, arch_analogy]
+    )
+
+    arch_next_btn.click(
+        page_architecture_next,
+        inputs=[arch_dropdown, layers_slider],
+        outputs=[page_arch_ui, page_hyper_ui, arch_analogy]
     )
 
-    gr.Markdown("### 2. Training Setup")
-
-    with gr.Row():
-        with gr.Column():
-            api_key = gr.Textbox(
-                label="Hugging Face Hub API Key",
-                type="password",
-                placeholder="hf_...",
-                info="Required only if pushing to hub."
-            )
-            repo_name = gr.Textbox(
-                label="Repository Name",
-                placeholder="your-username/your-model-name",
-                info="Required only if pushing to hub."
-            )
-            push_to_hub = gr.Checkbox(
-                label="Push to Hugging Face Hub?",
-                value=False
-            )
-
-    train_btn = gr.Button("🚀 Start Training", variant="primary")
-    output = gr.Textbox(label="Training Output", placeholder="Training logs and results will appear here...", lines=10)
-
-    train_btn.click(
-        fn=train_model,
-        inputs=[architecture_size, api_key, repo_name, push_to_hub],
-        outputs=output
+    hyper_next_btn.click(
+        page_hyperparams_next,
+        inputs=[lr_input, epochs_input, batch_input],
+        outputs=[page_hyper_ui, page_train_ui, train_status]
     )
 
-# Launch the app
-if __name__ == "__main__":
-    demo.launch()
+demo.queue().launch()