Update backend.py

backend.py

@@ -1,16 +1,17 @@
-# backend.py — REAL, FULL, WORKING VERSION
+# backend.py — REAL, FULL, WORKING VERSION — USER WRITES README, NO AUTO-GENERATED
 import sqlite3
 import threading
 import time
 import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader, Dataset
-from huggingface_hub import whoami
+from huggingface_hub import whoami, HfApi, create_repo
 from datasets import load_dataset
 from transformers import AutoTokenizer
 import psutil
 import os
 import signal
+import shutil
 
 DB_PATH = "llm_kitchen.db"
 training_queue = []
@@ -282,15 +283,22 @@ def run_training_job(job):
         device = "cuda" if torch.cuda.is_available() else "cpu"
         log_update(f"Run {run_id}: 🚀 Device = {device} | RAM available: {psutil.virtual_memory().available / (1024**3):.2f} GB", run_id)
 
+        # Load and save tokenizer
         tokenizer = AutoTokenizer.from_pretrained("gpt2")
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
 
+        tokenizer_save_path = f"./runs/{run_id}/tokenizer"
+        os.makedirs(tokenizer_save_path, exist_ok=True)
+        tokenizer.save_pretrained(tokenizer_save_path)
+        log_update(f"Run {run_id}: 💾 Tokenizer saved to {tokenizer_save_path}", run_id)
+
         vocab_size = len(tokenizer)
         model = get_model(job["arch_type"], vocab_size, job["num_layers"]).to(device)
         log_update(f"Run {run_id}: 🧱 Model initialized: {job['arch_type']} x{job['num_layers']} layers", run_id)
 
-        dataset = load_dataset("voidful/reasoning_gemini_300k", split="train[:100000]")
+        # Load dataset — reduce for testing
+        dataset = load_dataset("voidful/reasoning_gemini_300k", split="train[:5000]")
         def tokenize_function(examples):
             texts = [q + " " + a for q, a in zip(examples["message"], examples["answer"])]
             return tokenizer(texts, truncation=True, padding="max_length", max_length=128)
@@ -318,16 +326,19 @@ def run_training_job(job):
                 total_loss += loss.item()
                 if step % 50 == 0:
                     ram_gb = psutil.virtual_memory().used / (1024**3)
-                    log_update(f"Run {run_id}: Epoch {epoch+1} | Step {step} | Loss: {loss.item():.4f} | RAM: {ram_gb:.2f}GB", run_id)
+                    cpu_pct = psutil.cpu_percent(interval=0.1)
+                    log_update(f"Run {run_id}: Epoch {epoch+1} | Step {step} | Loss: {loss.item():.4f} | RAM: {ram_gb:.2f}GB | CPU: {cpu_pct}%", run_id)
 
             avg_loss = total_loss / len(train_loader)
             log_update(f"Run {run_id}: ✅ Epoch {epoch+1} completed | Avg Loss: {avg_loss:.4f}", run_id)
 
+        # Save model
         model_path = f"./runs/{run_id}"
         os.makedirs(model_path, exist_ok=True)
         torch.save(model.state_dict(), f"{model_path}/model.pth")
+        torch.save(model.state_dict(), f"{model_path}/pytorch_model.bin")  # HF standard
         update_run_status(run_id, "completed", f"Model saved to {model_path}")
-        log_update(f"Run {run_id}: 💾 Model checkpoint saved.", run_id)
+        log_update(f"Run {run_id}: 💾 Model checkpoint saved as .pth and .bin", run_id)
 
     except Exception as e:
         log_update(f"Run {run_id}: 💥 FAILED - {str(e)}", run_id)
@@ -345,6 +356,96 @@ def log_update(message, run_id):
     if run_id > 0:
         update_run_status(run_id, "running", full_msg)
 
+# ------------------------------ INFERENCE ------------------------------
+
+def run_inference(run_id, prompt):
+    """Load model + tokenizer and generate answer"""
+    model_path = f"./runs/{run_id}/model.pth"
+    tokenizer_path = f"./runs/{run_id}/tokenizer"
+
+    if not os.path.exists(model_path):
+        return "ModelError: Model not found. Did training complete?"
+
+    if not os.path.exists(tokenizer_path):
+        return "ModelError: Tokenizer not found."
+
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Fetch arch_type and num_layers from DB
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    cursor.execute("SELECT arch_type, num_layers FROM training_runs WHERE id = ?", (run_id,))
+    row = cursor.fetchone()
+    conn.close()
+
+    if not row:
+        return "ModelError: Run not found."
+
+    arch_type, num_layers = row
+    vocab_size = len(tokenizer)
+    model = get_model(arch_type, vocab_size, num_layers)
+
+    # Load weights
+    model.load_state_dict(torch.load(model_path, map_location="cpu"))
+    model.eval()
+
+    # Encode prompt
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=128)
+    input_ids = inputs.input_ids
+
+    # Generate (simple greedy)
+    with torch.no_grad():
+        outputs = model(input_ids)
+        logits = outputs["logits"]
+        predicted_token_ids = torch.argmax(logits, dim=-1)
+        generated_text = tokenizer.decode(predicted_token_ids[0], skip_special_tokens=True)
+
+    return f"🧑‍🍳 Model says:\n{generated_text}"
+
+# ------------------------------ PUBLISH TO HUB (USER WRITES README) ------------------------------
+
+def publish_run_to_hub(run_id, hf_token, repo_name, user_description=""):
+    """Save model + tokenizer and push to HF Hub. User provides description for README.md."""
+    local_dir = f"./runs/{run_id}/hub_upload"
+    os.makedirs(local_dir, exist_ok=True)
+
+    # Copy model as pytorch_model.bin
+    model_src = f"./runs/{run_id}/pytorch_model.bin"
+    if not os.path.exists(model_src):
+        raise FileNotFoundError("Model .bin not found for publishing.")
+    shutil.copy(model_src, f"{local_dir}/pytorch_model.bin")
+
+    # Copy tokenizer
+    tokenizer_path = f"./runs/{run_id}/tokenizer"
+    if not os.path.exists(tokenizer_path):
+        raise FileNotFoundError("Tokenizer not found.")
+    shutil.copytree(tokenizer_path, f"{local_dir}/tokenizer", dirs_exist_ok=True)
+
+    # Save user-provided description as README.md (or minimal fallback)
+    readme_content = user_description.strip()
+    if not readme_content:
+        readme_content = f"# Model from LLM Kitchen - Run #{run_id}\n\nTrained using custom architecture."
+
+    with open(f"{local_dir}/README.md", "w") as f:
+        f.write(readme_content)
+
+    # Create repo on user's account
+    api = HfApi()
+    repo_id = repo_name
+    url = create_repo(repo_id, token=hf_token, private=False, exist_ok=True)
+
+    # Push folder
+    api.upload_folder(
+        folder_path=local_dir,
+        repo_id=repo_id,
+        token=hf_token
+    )
+
+    return f"https://huggingface.co/{repo_id}"
+
 # ------------------------------ PUBLIC API ------------------------------
 
 __all__ = [
@@ -355,5 +456,7 @@ __all__ = [
     "start_training_if_free",
     "update_run_status",
     "log_update",
-    "init_db"
+    "init_db",
+    "run_inference",
+    "publish_run_to_hub"
 ]