Update backend.py

backend.py

@@ -1,17 +1,23 @@
-# backend.py
+# backend.py — REAL VERSION
 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 datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
+from transformers import AutoTokenizer
+import psutil
 import os
+import signal
 
 DB_PATH = "llm_kitchen.db"
 training_queue = []
 active_run_lock = threading.Lock()
 active_run_id = None
+RUN_TIMEOUT = 48 * 3600  # 48 hours
+MAX_RAM_PER_RUN_GB = 1.5
 
 # ------------------------------ DATABASE ------------------------------
 
@@ -38,6 +44,7 @@ def init_db():
             logs TEXT DEFAULT '',
             started_at DATETIME,
             completed_at DATETIME,
+            model_path TEXT,
             FOREIGN KEY (user_id) REFERENCES users(id)
         );
     """)
@@ -46,93 +53,20 @@ def init_db():
 
 init_db()
 
-def get_user_by_token(hf_token):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    cursor.execute("SELECT id FROM users WHERE hf_token = ?", (hf_token,))
-    row = cursor.fetchone()
-    conn.close()
-    return row[0] if row else None
-
-def create_user(hf_token):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    cursor.execute("INSERT INTO users (hf_token) VALUES (?)", (hf_token,))
-    user_id = cursor.lastrowid
-    conn.commit()
-    conn.close()
-    return user_id
-
-def create_training_run(user_id, config):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    cursor.execute("""
-        INSERT INTO training_runs 
-        (user_id, arch_type, num_layers, learning_rate, epochs, batch_size)
-        VALUES (?, ?, ?, ?, ?, ?)
-    """, (
-        user_id,
-        config['arch_type'],
-        config['num_layers'],
-        config['learning_rate'],
-        config['epochs'],
-        config['batch_size']
-    ))
-    run_id = cursor.lastrowid
-    conn.commit()
-    conn.close()
-    return run_id
-
-def get_user_runs(user_id):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    cursor.execute("""
-        SELECT id, arch_type, num_layers, status, started_at 
-        FROM training_runs 
-        WHERE user_id = ? 
-        ORDER BY started_at DESC
-    """, (user_id,))
-    runs = cursor.fetchall()
-    conn.close()
-    return runs
-
-def get_run_logs(run_id):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    cursor.execute("SELECT logs, status FROM training_runs WHERE id = ?", (run_id,))
-    row = cursor.fetchone()
-    conn.close()
-    return row if row else ("", "unknown")
-
-def update_run_status(run_id, status, logs=""):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    if status == 'running':
-        cursor.execute("UPDATE training_runs SET status = ?, started_at = CURRENT_TIMESTAMP WHERE id = ?", (status, run_id))
-    elif status in ['completed', 'failed', 'timeout']:
-        cursor.execute("UPDATE training_runs SET status = ?, completed_at = CURRENT_TIMESTAMP WHERE id = ?", (status, run_id))
-    if logs:
-        current_logs = get_run_logs(run_id)[0]
-        cursor.execute("UPDATE training_runs SET logs = ? WHERE id = ?", (current_logs + "\n" + logs, run_id))
-    conn.commit()
-    conn.close()
+# ... [KEEP ALL DB HELPER FUNCTIONS: get_user_by_token, create_user, etc. — NO CHANGES] ...
 
 # ------------------------------ AUTH ------------------------------
-
-def verify_hf_token(token):
-    try:
-        whoami(token=token)
-        user_id = get_user_by_token(token)
-        if not user_id:
-            user_id = create_user(token)
-            return user_id, "Welcome to the LLM Kitchen, Chef! 🍳 Your apron is ready."
-        else:
-            return user_id, "Welcome back, Chef! 👨‍🍳 Your last dish is still warm."
-    except Exception as e:
-        return None, f"Invalid token. Please try again. ({str(e)})"
+# ... [KEEP verify_hf_token — NO CHANGES] ...
 
 # ------------------------------ TRAINING QUEUE ------------------------------
 
+def ram_available():
+    """Check if we can start a new run (1.5GB per run)"""
+    total_ram = psutil.virtual_memory().total / (1024**3)  # GB
+    used_ram = psutil.virtual_memory().used / (1024**3)    # GB
+    available_gb = total_ram - used_ram
+    return available_gb >= MAX_RAM_PER_RUN_GB
+
 def queue_training_run(user_id, config):
     run_id = create_training_run(user_id, config)
     training_queue.append({
@@ -142,11 +76,6 @@ def queue_training_run(user_id, config):
     })
     return run_id
 
-def ram_check_mock():
-    # Mock: Allow 1 run at a time, 1.5GB per run
-    global active_run_id
-    return active_run_id is None
-
 def start_training_if_free():
     global active_run_id
     with active_run_lock:
@@ -154,7 +83,8 @@ def start_training_if_free():
             return False
         if not training_queue:
             return False
-        if not ram_check_mock():
+        if not ram_available():
+            log_update("MemoryWarning: Not enough RAM to start new run.", -1)
             return False
 
         job = training_queue.pop(0)
@@ -163,72 +93,184 @@ def start_training_if_free():
 
         thread = threading.Thread(target=run_training_job, args=(job,))
         thread.start()
+
+        # Start 48h timeout killer
+        timer = threading.Timer(RUN_TIMEOUT, kill_run_timeout, args=[job["run_id"]])
+        timer.start()
+
         return True
 
+def kill_run_timeout(run_id):
+    global active_run_id
+    with active_run_lock:
+        if active_run_id == run_id:
+            log_update(f"Run {run_id}: 💥 48-HOUR TIMEOUT REACHED. Terminating.", run_id)
+            update_run_status(run_id, "timeout")
+            active_run_id = None
+            start_training_if_free()  # try next
+
+# ------------------------------ CUSTOM MODELS FROM SCRATCH ------------------------------
+
+class CNNLanguageModel(nn.Module):
+    def __init__(self, vocab_size, embed_dim=128, num_layers=4):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        layers = []
+        in_ch = embed_dim
+        for _ in range(num_layers):
+            layers.append(nn.Conv1d(in_ch, in_ch * 2, kernel_size=3, padding=1))
+            layers.append(nn.ReLU())
+            in_ch *= 2
+        self.convs = nn.Sequential(*layers)
+        self.fc = nn.Linear(in_ch, vocab_size)
+
+    def forward(self, x, labels=None):
+        x = self.embedding(x).transpose(1, 2)  # (B, E, L)
+        x = self.convs(x).transpose(1, 2)     # (B, L, E*2^N)
+        logits = self.fc(x)
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+        return {"loss": loss, "logits": logits}
+
+class RNNLanguageModel(nn.Module):
+    def __init__(self, vocab_size, embed_dim=128, hidden_dim=256, num_layers=2):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.rnn = nn.LSTM(embed_dim, hidden_dim, num_layers, batch_first=True)
+        self.fc = nn.Linear(hidden_dim, vocab_size)
+
+    def forward(self, x, labels=None):
+        x = self.embedding(x)
+        output, _ = self.rnn(x)
+        logits = self.fc(output)
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+        return {"loss": loss, "logits": logits}
+
+class TransformerLanguageModel(nn.Module):
+    def __init__(self, vocab_size, embed_dim=128, num_heads=4, num_layers=3):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, batch_first=True)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.fc = nn.Linear(embed_dim, vocab_size)
+
+    def forward(self, x, labels=None):
+        x = self.embedding(x)
+        x = self.transformer(x)
+        logits = self.fc(x)
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+        return {"loss": loss, "logits": logits}
+
+def get_model(arch_type, vocab_size, num_layers):
+    if arch_type == "cnn":
+        return CNNLanguageModel(vocab_size, num_layers=num_layers)
+    elif arch_type == "rnn":
+        return RNNLanguageModel(vocab_size, num_layers=num_layers)
+    elif arch_type == "transformer":
+        return TransformerLanguageModel(vocab_size, num_layers=num_layers)
+    else:
+        raise ValueError(f"Unknown arch: {arch_type}")
+
+# ------------------------------ DATASET ------------------------------
+
+class TextDataset(Dataset):
+    def __init__(self, tokenized_data):
+        self.input_ids = tokenized_data["input_ids"]
+        self.labels = tokenized_data["input_ids"]  # causal LM
+
+    def __len__(self):
+        return len(self.input_ids)
+
+    def __getitem__(self, idx):
+        return {
+            "input_ids": torch.tensor(self.input_ids[idx], dtype=torch.long),
+            "labels": torch.tensor(self.labels[idx], dtype=torch.long),
+        }
+
+# ------------------------------ TRAINING JOB ------------------------------
+
 def run_training_job(job):
     global active_run_id
     run_id = job["run_id"]
     try:
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        log_update(f"Run {run_id}: Device = {device}", run_id)
+        log_update(f"Run {run_id}: 🚀 Device = {device} | RAM available: {psutil.virtual_memory().available / (1024**3):.2f} GB", run_id)
 
-        # Load tiny model for demo (replace with custom later)
-        model_name = "distilgpt2"
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        # Load tokenizer (shared for all models)
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
 
-        model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
-        log_update(f"Run {run_id}: Model loaded", run_id)
+        vocab_size = len(tokenizer)
+
+        # Build model from scratch
+        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)
 
-        # Load dataset
-        dataset = load_dataset("voidful/reasoning_gemini_300k", split="train[:1%]")  # Tiny slice for demo
+        # Load dataset — full training set (or 100K for speed)
+        dataset = load_dataset("voidful/reasoning_gemini_300k", split="train[:100000]")
         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)
-        tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["message", "answer"])
-        log_update(f"Run {run_id}: Dataset tokenized", run_id)
-
-        # Training args
-        training_args = TrainingArguments(
-            output_dir=f"./runs/{run_id}",
-            num_train_epochs=job["epochs"],
-            per_device_train_batch_size=job["batch_size"],
-            learning_rate=job["learning_rate"],
-            save_strategy="no",
-            logging_steps=1,
-            report_to="none",
-            fp16=False,
-            no_cuda=(device == "cpu")
-        )
-
-        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=tokenized_dataset,
-            data_collator=data_collator,
-        )
-
-        log_update(f"Run {run_id}: Starting training...", run_id)
-        trainer.train()
-
-        # Simulate 48h timeout with short sleep for demo
-        time.sleep(10)  # Replace with real training
-
-        eval_results = trainer.evaluate()
-        log_update(f"Run {run_id}: Training complete. Loss = {eval_results.get('eval_loss', 'N/A')}", run_id)
-        update_run_status(run_id, "completed")
+        tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names)
+        train_dataset = TextDataset(tokenized_dataset)
+
+        train_loader = DataLoader(train_dataset, batch_size=job["batch_size"], shuffle=True)
+
+        # Optimizer
+        optimizer = torch.optim.AdamW(model.parameters(), lr=job["learning_rate"])
+
+        # Training loop
+        model.train()
+        log_update(f"Run {run_id}: ▶️ Starting training for {job['epochs']} epochs...", run_id)
+
+        for epoch in range(job["epochs"]):
+            total_loss = 0
+            for step, batch in enumerate(train_loader):
+                input_ids = batch["input_ids"].to(device)
+                labels = batch["labels"].to(device)
+
+                optimizer.zero_grad()
+                outputs = model(input_ids, labels=labels)
+                loss = outputs["loss"]
+                loss.backward()
+                optimizer.step()
+
+                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)
+
+            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")
+        update_run_status(run_id, "completed", f"Model saved to {model_path}")
+        log_update(f"Run {run_id}: 💾 Model checkpoint saved.", run_id)
 
     except Exception as e:
-        log_update(f"Run {run_id}: FAILED - {str(e)}", run_id)
+        log_update(f"Run {run_id}: 💥 FAILED - {str(e)}", run_id)
         update_run_status(run_id, "failed")
     finally:
         with active_run_lock:
-            active_run_id = None
-        # Try starting next queued job
+            if active_run_id == run_id:
+                active_run_id = None
         start_training_if_free()
 
 def log_update(message, run_id):
-    print(f"[LOG] {message}")  # Also print to Spaces logs
-    update_run_status(run_id, "running", message)
+    timestamp = time.strftime("%H:%M:%S")
+    full_msg = f"[{timestamp}] {message}"
+    print(full_msg)  # Also shows in HF Spaces logs
+    if run_id > 0:
+        update_run_status(run_id, "running", full_msg)