Upload app.py with huggingface_hub

app.py

@@ -1,1036 +1,690 @@
 #!/usr/bin/env python3
 """
-OpenLLM Training Space - Main Application
-
-This is the main entry point for the Hugging Face Space.
-It provides a web interface for running OpenLLM training with authentication.
-
-Author: Louis Chua Bean Chong
-License: GPLv3
+OpenLLM Real Models App - Final working version with correct attribute naming
 """
 
-import os
-import sys
-from pathlib import Path
-
 import gradio as gr
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import json
+import logging
+import sentencepiece as spm
+import math
+from pathlib import Path
+from typing import Dict, Any, Optional
+from huggingface_hub import snapshot_download
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class GPTConfig:
+    """GPT model configuration"""
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        n_layer=6,
+        n_head=8,
+        n_embd=512,
+        block_size=1024,
+        dropout=0.1,
+        bias=False,
+        **kwargs,
+    ):
+        # Accept any additional kwargs to handle extra config fields
+        self.vocab_size = vocab_size
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_embd = n_embd
+        self.block_size = block_size
+        self.dropout = dropout
+        self.bias = bias
+
+
+class GPT(nn.Module):
+    """GPT-style transformer model - EXACT architecture matching the saved model"""
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+
+        # Create the transformer module with the exact naming convention
+        self.transformer = nn.ModuleDict(
+            dict(
+                wte=nn.Embedding(config.vocab_size, config.n_embd),
+                wpe=nn.Embedding(config.block_size, config.n_embd),
+                drop=nn.Dropout(config.dropout),
+                h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+                ln_f=nn.LayerNorm(config.n_embd),
+            )
+        )
 
-# Add core/src to path for imports
-core_src_path = str(Path(__file__).parent / "core" / "src")
-if core_src_path not in sys.path:
-    sys.path.insert(0, core_src_path)
-
-# Import our authentication and training modules
-try:
-    from openllm_training_with_auth import OpenLLMTrainingManager
-    from space_auth_test import test_space_authentication
-
-    MODULES_AVAILABLE = True
-except ImportError as e:
-    MODULES_AVAILABLE = False
-    print(f"❌ Required modules not available: {e}")
-
+        # Language model head - Use bias=False to match saved models
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        # Initialize weights
+        self.apply(self._init_weights)
+        for pn, p in self.named_parameters():
+            if pn.endswith("c_proj.weight"):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer))
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        device = idx.device
+        b, t = idx.size()
+        assert (
+            t <= self.config.block_size
+        ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+
+        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)
+        tok_emb = self.transformer.wte(idx)
+        pos_emb = self.transformer.wpe(pos)
+        x = self.transformer.drop(tok_emb + pos_emb)
+
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+
+        if targets is not None:
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1
+            )
+        else:
+            logits = self.lm_head(x[:, [-1], :])
+            loss = None
+
+        return logits, loss
+
+    def generate(
+        self, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None, do_sample=True
+    ):
+        for _ in range(max_new_tokens):
+            idx_cond = (
+                idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size :]
+            )
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("Inf")
+
+            if top_p is not None:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                sorted_indices_to_remove[..., 0] = 0
+                indices_to_remove = sorted_indices_to_remove.scatter(
+                    1, sorted_indices, sorted_indices_to_remove
+                )
+                logits[indices_to_remove] = -float("Inf")
 
-def create_space_interface():
-    """Create the Gradio interface for the Space."""
+            probs = F.softmax(logits, dim=-1)
+            if do_sample:
+                idx_next = torch.multinomial(probs, num_samples=1)
+            else:
+                _, idx_next = torch.topk(probs, k=1, dim=-1)
+
+            idx = torch.cat((idx, idx_next), dim=1)
+
+        return idx
+
+
+class Block(nn.Module):
+    """Transformer block with self-attention and feed-forward layers"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class CausalSelfAttention(nn.Module):
+    """Multi-head self-attention with causal masking - FINAL WORKING VERSION"""
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        self.use_bias = config.bias  # Use different name for the boolean flag
+
+        # REGISTER THE ATTENTION BIAS as a buffer (not parameter) to match saved model
+        # This is actually an attention mask, not a learnable bias
+        if config.bias:
+            # Create a causal attention mask buffer
+            mask = torch.tril(torch.ones(config.block_size, config.block_size))
+            mask = mask.view(1, 1, config.block_size, config.block_size)
+            self.register_buffer("bias", mask)  # This matches the saved model's 'bias' key
+        else:
+            self.register_buffer("bias", None)
+
+    def forward(self, x):
+        B, T, C = x.size()
+
+        # Calculate query, key, values for all heads
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+
+        # Causal self-attention using the bias mask
+        if self.bias is not None:
+            # Use the causal mask
+            attn_mask = self.bias[:, :, :T, :T]
+            y = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=attn_mask,
+                dropout_p=self.dropout if self.training else 0,
+                is_causal=False,
+            )
+        else:
+            # Use built-in causal attention
+            y = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=None,
+                dropout_p=self.dropout if self.training else 0,
+                is_causal=True,
+            )
 
-    def run_authentication_test():
-        """Run the authentication test and return results."""
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+
+        # Output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+
+class MLP(nn.Module):
+    """Multi-layer perceptron"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu = nn.GELU()
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+
+class RealOpenLLMInference:
+    """Real OpenLLM inference engine using actual trained models"""
+
+    def __init__(self):
+        self.models = {}
+        self.tokenizers = {}
+        self.current_model = None
+
+        # Real model configurations from Hugging Face
+        self.model_configs = {
+            "openllm-small-extended-4k": {
+                "name": "OpenLLM Small (4k steps)",
+                "description": "Real model trained for 4,000 steps - Early training stage",
+                "hf_repo": "lemms/openllm-small-extended-4k",
+                "training_steps": 4000,
+                "parameters": "35.8M",
+            },
+            "openllm-small-extended-6k": {
+                "name": "OpenLLM Small (6k steps)",
+                "description": "Real model trained for 6,000 steps - Improved coherence (Perplexity: 816.040)",
+                "hf_repo": "lemms/openllm-small-extended-6k",
+                "training_steps": 6000,
+                "parameters": "35.8M",
+            },
+            "openllm-small-extended-7k": {
+                "name": "OpenLLM Small (7k steps)",
+                "description": "Real model trained for 7,000 steps - Enhanced quality (Loss: 2.100, Perplexity: 8.200)",
+                "hf_repo": "lemms/openllm-small-extended-7k",
+                "training_steps": 7000,
+                "parameters": "35.8M",
+            },
+            "openllm-small-extended-8k": {
+                "name": "OpenLLM Small (8k steps)",
+                "description": "Real model trained for 8,000 steps - Sophisticated understanding",
+                "hf_repo": "lemms/openllm-small-extended-8k",
+                "training_steps": 8000,
+                "parameters": "35.8M",
+            },
+            "openllm-small-extended-9k": {
+                "name": "OpenLLM Small (9k steps)",
+                "description": "Real model trained for 9,000 steps - Best performing model",
+                "hf_repo": "lemms/openllm-small-extended-9k",
+                "training_steps": 9000,
+                "parameters": "35.8M",
+            },
+            "openllm-small-extended-10k": {
+                "name": "OpenLLM Small (10k steps)",
+                "description": "Real model trained for 10,000 steps - Latest extended training",
+                "hf_repo": "lemms/openllm-small-extended-10k",
+                "training_steps": 10000,
+                "parameters": "35.8M",
+            },
+            "openllm-small-extended-10k-improved": {
+                "name": "OpenLLM Small (10k steps - Improved)",
+                "description": "Real model trained for 10,000 steps with improved training process - Proper checkpoint format",
+                "hf_repo": "lemms/openllm-small-extended-10k-improved",
+                "training_steps": 10000,
+                "parameters": "35.8M",
+            },
+        }
+
+        logger.info("🚀 Real OpenLLM Inference Engine initialized")
+
+    def load_model_from_hf(self, model_id: str) -> bool:
+        """Load a real model from Hugging Face"""
         try:
-            if not MODULES_AVAILABLE:
-                return "❌ Required modules not available. Please check deployment."
-
-            # Capture output from authentication test
-            import contextlib
-            import io
-
-            output = io.StringIO()
-            with contextlib.redirect_stdout(output):
-                success = test_space_authentication()
+            config = self.model_configs.get(model_id)
+            if not config:
+                logger.error(f"❌ Unknown model ID: {model_id}")
+                return False
+
+            logger.info(f"📥 Loading real model from HF: {config['hf_repo']}")
+
+            # Download model from Hugging Face
+            local_dir = snapshot_download(
+                repo_id=config["hf_repo"],
+                repo_type="model",
+                local_dir=f"temp_{model_id}",
+                allow_patterns=["*.pt", "*.json", "*.model", "*.bin"],
+            )
 
-            result = output.getvalue()
+            logger.info(f"✅ Downloaded model to: {local_dir}")
 
+            # Load model and tokenizer
+            success = self._load_model_and_tokenizer(local_dir, model_id)
             if success:
-                return f"✅ Authentication Test Results:\n\n{result}"
+                self.current_model = model_id
+                logger.info(f"✅ Successfully loaded real model: {model_id}")
+                return True
             else:
-                return f"❌ Authentication Test Failed:\n\n{result}"
+                return False
 
         except Exception as e:
-            return f"❌ Error running authentication test: {e}"
+            logger.error(f"❌ Failed to load real model from HF {model_id}: {e}")
+            return False
 
-    def run_training(model_size, training_steps, use_real_training=False):
-        """Run the OpenLLM training with authentication."""
+    def _load_model_and_tokenizer(self, model_dir: str, model_id: str) -> bool:
+        """Load model and tokenizer from local directory"""
         try:
-            if not MODULES_AVAILABLE:
-                return "❌ Required modules not available. Please check deployment."
-
-            # Security mitigation: Input validation and sanitization
-            if not isinstance(model_size, str) or model_size not in ["small", "medium", "large"]:
-                return "❌ Invalid model size. Must be 'small', 'medium', or 'large'."
-
-            if (
-                not isinstance(training_steps, (int, float))
-                or training_steps < 1000
-                or training_steps > 50000
-            ):
-                return "❌ Invalid training steps. Must be between 1000 and 50000."
-
-            # Sanitize inputs
-            model_size = str(model_size).strip().lower()
-            training_steps = int(float(training_steps))
-
-            print(f"🚀 Starting OpenLLM Training")
-            print("=" * 50)
-            print(f"📊 Model Size: {model_size}")
-            print(f"🔄 Training Steps: {training_steps}")
-            print(f"🎯 Training Mode: {'Real Training' if use_real_training else 'Demonstration'}")
-
-            if use_real_training:
-                # Use real training with comprehensive features
-                try:
-                    from real_training_manager import RealTrainingManager, TrainingConfig
-
-                    # Create configuration for real training
-                    config = TrainingConfig(
-                        model_size=model_size,
-                        training_steps=training_steps,
-                        batch_size=32 if model_size == "small" else 16,
-                        learning_rate=3e-4,
-                        data_file="data/clean/training_data.txt",
-                        save_every=1000,
-                        eval_every=500,
-                    )
-
-                    # Initialize real training manager
-                    manager = RealTrainingManager(config)
-
-                    # Run real training
-                    model = manager.train()
-
-                    # Upload model
-                    repo_id = manager.upload_model(model)
-
-                    if repo_id:
-                        result = f"✅ Real Training completed successfully!\n\n"
-                        result += f"📊 Results:\n"
-                        result += f"   - Model Size: {model_size}\n"
-                        result += f"   - Training Steps: {training_steps}\n"
-                        result += f"   - Final Loss: {manager.training_history[-1]['loss']:.4f}\n"
-                        result += f"   - Best Validation Loss: {manager.best_loss:.4f}\n"
-                        result += f"   - Model URL: https://huggingface.co/{repo_id}\n\n"
-                        result += f"🎉 Model available at: https://huggingface.co/{repo_id}"
-                    else:
-                        result = f"⚠️ Real training completed but upload failed\n\n"
-                        result += f"📊 Results:\n"
-                        result += f"   - Model Size: {model_size}\n"
-                        result += f"   - Training Steps: {training_steps}\n"
-                        result += f"   - Final Loss: {manager.training_history[-1]['loss']:.4f}\n"
-                        result += f"   - Model saved locally: ./trained_model"
-
-                    return result
-
-                except ImportError:
-                    return (
-                        "❌ Real training module not available. Falling back to demonstration mode."
-                    )
-                except Exception as e:
-                    return (
-                        f"❌ Real training failed: {str(e)}\n\nFalling back to demonstration mode."
-                    )
-
-            # Fallback to demonstration training
-            import contextlib
-            import io
-
-            output = io.StringIO()
-            with contextlib.redirect_stdout(output):
-                training_manager = OpenLLMTrainingManager()
-                repo_id = training_manager.run_training(model_size=model_size, steps=training_steps)
+            model_path = Path(model_dir)
 
-            result = output.getvalue()
+            # Load model configuration
+            config_file = model_path / "config.json"
+            if config_file.exists():
+                with open(config_file, "r") as f:
+                    config_data = json.load(f)
 
-            return f"✅ Training Results:\n\n{result}\n\n🎉 Model available at: https://huggingface.co/{repo_id}"
+                logger.info(f"📋 Config data keys: {list(config_data.keys())}")
 
-        except Exception as e:
-            return f"❌ Error running training: {e}"
-
-    def resume_training_from_7k_to_8k():
-        """Resume training from 7k model to create 8k model."""
-        try:
-            if not MODULES_AVAILABLE:
-                return "❌ Required modules not available. Please check deployment."
-
-            # Import required modules
-            import json
-            import time
-            from datetime import datetime
-
-            import torch
-            from huggingface_hub import HfApi, create_repo, snapshot_download, whoami
-            from model import GPTConfig, GPTModel
-            from train_model import TextDataLoader
-
-            print("🚀 Resuming Training from 7k to 8k Model")
-            print("=" * 50)
-
-            # Configuration
-            hf_model_id = "lemms/openllm-small-extended-7k"
-            additional_steps = 1000  # Train for 1000 more steps to reach 8k
-            total_steps = 8000  # Total steps for the new model
-
-            print(f"📥 Source Model: {hf_model_id}")
-            print(f"📈 Additional Steps: {additional_steps}")
-            print(f"🎯 Target Steps: {total_steps}")
-
-            # Setup authentication
-            print("🔐 Setting up authentication...")
-            try:
-                user_info = whoami()
-                username = user_info.get("name", "unknown")
-                print(f"✅ Authentication successful! User: {username}")
-            except Exception as e:
-                return f"❌ Authentication failed: {e}"
-
-            # Setup device
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            print(f"🖥️ Using device: {device}")
-
-            # Load model from Hugging Face
-            print(f"📥 Loading model from Hugging Face: {hf_model_id}")
-            try:
-                local_dir = snapshot_download(
-                    repo_id=hf_model_id,
-                    repo_type="model",
-                    local_dir=f"downloaded_models/{hf_model_id.replace('/', '_')}",
-                )
-                print(f"✅ Model downloaded to: {local_dir}")
-
-                # Load config
-                config_path = Path(local_dir) / "config.json"
-                if config_path.exists():
-                    with open(config_path, "r") as f:
-                        config_data = json.load(f)
-
-                    config = GPTConfig(
-                        vocab_size=config_data.get("vocab_size", 32000),
-                        block_size=config_data.get("block_size", 1024),
-                        n_layer=config_data.get("n_layer", 6),
-                        n_head=config_data.get("n_head", 6),
-                        n_embd=config_data.get("n_embd", 384),
-                    )
-                    print(f"📊 Loaded model config: {config}")
-                else:
-                    config = GPTConfig.small()
-                    config.vocab_size = 32000
-                    print(f"⚠️ Config file not found, using default config")
-
-                # Create model and load weights
-                model = GPTModel(config)
-                model_path = Path(local_dir) / "pytorch_model.bin"
-
-                if model_path.exists():
-                    state_dict = torch.load(model_path, map_location=device)
-                    model.load_state_dict(state_dict)
-                    print(f"✅ Model weights loaded successfully")
+                # Handle different config structures
+                if "model_config" in config_data:
+                    # Extract model_config section
+                    model_config_data = config_data["model_config"]
                 else:
-                    raise FileNotFoundError(f"Model file not found: {model_path}")
-
-                model = model.to(device)
-
-            except Exception as e:
-                return f"❌ Failed to load model from Hugging Face: {e}"
-
-            # Create data loaders
-            print(f"📊 Loading training data...")
-            tokenizer_path = "data/tokenizer/tokenizer.model"
-
-            train_loader = TextDataLoader(
-                data_file="data/clean/training_data.txt",
-                tokenizer_path=tokenizer_path,
-                seq_len=1024,
-                batch_size=16,
-                shuffle=True,
-            )
-
-            print(f"✅ Data loader created")
-
-            # Setup optimizer
-            optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
-
-            # Training loop
-            print(f"\n🔄 Starting training loop...")
-            start_time = time.time()
-            training_history = []
-            best_loss = float("inf")
-
-            try:
-                train_iterator = iter(train_loader)
-
-                for step in range(additional_steps):
-                    # Get batch
-                    try:
-                        batch = next(train_iterator)
-                    except StopIteration:
-                        # Restart data loader if exhausted
-                        train_loader = TextDataLoader(
-                            data_file="data/clean/training_data.txt",
-                            tokenizer_path=tokenizer_path,
-                            seq_len=1024,
-                            batch_size=16,
-                            shuffle=True,
-                        )
-                        train_iterator = iter(train_loader)
-                        batch = next(train_iterator)
-
-                    # Prepare inputs
-                    if isinstance(batch, (list, tuple)):
-                        inputs = batch[0].to(device)
-                        targets = batch[1].to(device) if len(batch) > 1 else None
-                    else:
-                        inputs = batch.to(device)
-                        targets = None
-
-                    # Forward pass
-                    logits, loss = model(inputs, targets)
-
-                    # Backward pass
-                    loss.backward()
-
-                    # Optimizer step
-                    optimizer.step()
-                    optimizer.zero_grad()
-
-                    # Record training history
-                    training_history.append(
-                        {
-                            "step": 7000 + step + 1,  # Continue from step 7000
-                            "loss": loss.item(),
-                            "timestamp": datetime.now().isoformat(),
-                        }
-                    )
-
-                    # Progress reporting
-                    if (step + 1) % 10 == 0:
-                        elapsed = time.time() - start_time
-                        steps_per_sec = (step + 1) / elapsed
-                        eta = (additional_steps - step - 1) / steps_per_sec
-
-                        print(
-                            f"Step {7000 + step + 1}/{total_steps} | "
-                            f"Loss: {loss.item():.4f} | "
-                            f"Speed: {steps_per_sec:.1f} steps/s | "
-                            f"ETA: {eta/60:.1f} min"
-                        )
-
-                    # Evaluation
-                    if (step + 1) % 250 == 0:
-                        model.eval()
-                        total_loss = 0.0
-                        num_batches = 0
-
-                        with torch.no_grad():
-                            for val_batch in train_loader:  # Use same loader for simplicity
-                                if isinstance(val_batch, (list, tuple)):
-                                    val_inputs = val_batch[0].to(device)
-                                    val_targets = (
-                                        val_batch[1].to(device) if len(val_batch) > 1 else None
-                                    )
-                                else:
-                                    val_inputs = val_batch.to(device)
-                                    val_targets = None
-
-                                val_logits, val_loss = model(val_inputs, val_targets)
-                                total_loss += val_loss.item()
-                                num_batches += 1
-
-                                if num_batches >= 5:  # Limit evaluation
-                                    break
-
-                        avg_val_loss = total_loss / num_batches
-                        model.train()
-                        print(f"📊 Validation Loss: {avg_val_loss:.4f}")
-
-                        # Check for best model
-                        if avg_val_loss < best_loss:
-                            best_loss = avg_val_loss
-                            print(f"🏆 New best validation loss: {best_loss:.4f}")
-
-                print(f"\n🎉 Training completed successfully!")
-                print(f"📊 Final Results:")
-                print(f"   - Additional Steps: {additional_steps}")
-                print(f"   - Total Steps: {total_steps}")
-                print(f"   - Final Loss: {loss.item():.4f}")
-                print(f"   - Best Validation Loss: {best_loss:.4f}")
-                print(f"   - Training Time: {(time.time() - start_time)/3600:.2f} hours")
-
-                # Upload model
-                print(f"\n📤 Uploading model to Hugging Face Hub...")
-
-                # Create model directory
-                model_path = Path("./trained_model")
-                model_path.mkdir(exist_ok=True)
-
-                # Save model files
-                torch.save(model.state_dict(), model_path / "pytorch_model.bin")
-
-                # Save config
-                config_dict = {
-                    "model_type": "openllm",
-                    "model_size": "small",
-                    "vocab_size": 32000,
-                    "block_size": 1024,
-                    "n_layer": 6,
-                    "n_head": 6,
-                    "n_embd": 384,
-                    "training_config": {
-                        "model_size": "small",
-                        "training_steps": total_steps,
-                        "additional_steps": additional_steps,
-                        "base_model": hf_model_id,
-                    },
-                    "training_history": training_history,
+                    # Use the entire config as model config
+                    model_config_data = config_data
+
+                # Create GPTConfig with only the expected parameters
+                expected_params = {
+                    "vocab_size",
+                    "n_layer",
+                    "n_head",
+                    "n_embd",
+                    "block_size",
+                    "dropout",
+                    "bias",
                 }
 
-                with open(model_path / "config.json", "w") as f:
-                    json.dump(config_dict, f, indent=2)
-
-                # Create model card
-                readme_content = f"""# OpenLLM Small Model - Extended to 8k Steps
-
-This is an OpenLLM small model trained for {total_steps} steps by resuming training from [lemms/openllm-small-extended-7k](https://huggingface.co/lemms/openllm-small-extended-7k).
-
-## Model Details
-
-- **Model Type**: OpenLLM
-- **Size**: small
-- **Training Steps**: {total_steps}
-- **Additional Steps**: {additional_steps}
-- **Base Model**: [lemms/openllm-small-extended-7k](https://huggingface.co/lemms/openllm-small-extended-7k)
-- **Final Loss**: {training_history[-1]['loss']:.4f} if training_history else 'N/A'
-- **Framework**: PyTorch
-- **License**: GPL-3.0
-
-## Training Configuration
-
-```json
-{json.dumps(config_dict, indent=2)}
-```
-
-## Training History
-
-The model was trained with the following key metrics:
-- Best validation loss: {best_loss:.4f}
-- Total training time: {len(training_history)} steps
-- Device used: {device}
-
-## Usage
-
-This model can be used for text generation and language modeling tasks.
-
-## Author
-
-Louis Chua Bean Chong
-
-## License
-
-GPL-3.0
-"""
-
-                with open(model_path / "README.md", "w") as f:
-                    f.write(readme_content)
-
-                # Upload to Hugging Face
-                repo_name = "openllm-small-extended-8k"
-                repo_id = f"{username}/{repo_name}"
-
-                try:
-                    # Create repository
-                    create_repo(repo_id=repo_id, repo_type="model", exist_ok=True, private=False)
-
-                    # Upload files
-                    api = HfApi()
-                    api.upload_folder(
-                        folder_path=str(model_path),
-                        repo_id=repo_id,
-                        repo_type="model",
-                        commit_message=f"Add OpenLLM small model extended to {total_steps} steps",
-                    )
-
-                    print(f"✅ Model uploaded successfully!")
-                    print(f"🔗 Model URL: https://huggingface.co/{repo_id}")
-
-                    result = f"✅ 8k Model Training completed successfully!\n\n"
-                    result += f"📊 Results:\n"
-                    result += f"   - Base Model: {hf_model_id}\n"
-                    result += f"   - Additional Steps: {additional_steps}\n"
-                    result += f"   - Total Steps: {total_steps}\n"
-                    result += f"   - Final Loss: {loss.item():.4f}\n"
-                    result += f"   - Best Validation Loss: {best_loss:.4f}\n"
-                    result += f"   - Model URL: https://huggingface.co/{repo_id}\n\n"
-                    result += f"🎉 Extended model available at: https://huggingface.co/{repo_id}"
+                config_kwargs = {}
+                for key, value in model_config_data.items():
+                    if key in expected_params:
+                        config_kwargs[key] = value
 
-                    return result
+                logger.info(f"🔧 Using config parameters: {config_kwargs}")
+                model_config = GPTConfig(**config_kwargs)
+            else:
+                # Default configuration for OpenLLM small models
+                model_config = GPTConfig(
+                    vocab_size=32000,
+                    n_layer=6,
+                    n_head=8,
+                    n_embd=512,
+                    block_size=1024,
+                    dropout=0.1,
+                    bias=False,
+                )
 
-                except Exception as e:
-                    print(f"❌ Model upload failed: {e}")
-                    return f"⚠️ Training completed but upload failed: {e}"
+            # Load model weights
+            model_file = model_path / "best_model.pt"
+            if not model_file.exists():
+                model_file = model_path / "model.pt"
+            if not model_file.exists():
+                model_file = model_path / "pytorch_model.bin"
+
+            if model_file.exists():
+                logger.info(f"📦 Loading model from: {model_file}")
+                model = GPT(model_config)
+                checkpoint = torch.load(model_file, map_location="cpu")
+
+                # Handle different checkpoint formats
+                if isinstance(checkpoint, dict):
+                    if "model_state_dict" in checkpoint:
+                        # Extract the actual model weights
+                        state_dict = checkpoint["model_state_dict"]
+                        logger.info(f"📋 Loading from model_state_dict with {len(state_dict)} keys")
+                    elif "model" in checkpoint:
+                        state_dict = checkpoint["model"]
+                        logger.info(f"📋 Loading from model with {len(state_dict)} keys")
+                    else:
+                        # Try to load directly as state dict
+                        state_dict = checkpoint
+                        logger.info(f"📋 Loading direct state dict with {len(state_dict)} keys")
+                else:
+                    # Direct state dict
+                    state_dict = checkpoint
+                    logger.info(f"📋 Loading direct state dict with {len(state_dict)} keys")
+
+                # Load the state dict
+                model.load_state_dict(state_dict)
+                model.eval()
+                self.models[model_id] = model
+                logger.info(f"✅ Model loaded successfully")
+            else:
+                logger.error(f"❌ Model file not found in {model_dir}")
+                logger.error(f"   Available files: {list(model_path.glob('*'))}")
+                return False
+
+            # Load tokenizer
+            tokenizer_file = model_path / "tokenizer.model"
+            if tokenizer_file.exists():
+                tokenizer = spm.SentencePieceProcessor()
+                tokenizer.load(str(tokenizer_file))
+                self.tokenizers[model_id] = tokenizer
+                logger.info(f"✅ Tokenizer loaded successfully")
+            else:
+                logger.error(f"❌ Tokenizer file not found in {model_dir}")
+                return False
 
-            except KeyboardInterrupt:
-                print(f"\n⚠️ Training interrupted by user")
-                return "⚠️ Training was interrupted by user"
+            return True
 
         except Exception as e:
-            return f"❌ Error resuming training: {e}"
+            logger.error(f"❌ Failed to load model and tokenizer: {e}")
+            import traceback
+
+            logger.error(f"📋 Full traceback: {traceback.format_exc()}")
+            return False
+
+    def generate_text(
+        self,
+        prompt: str,
+        max_length: int = 100,
+        temperature: float = 0.7,
+        top_k: int = 50,
+        top_p: float = 0.9,
+    ) -> str:
+        """Generate text using the loaded real model"""
+        if not self.current_model or self.current_model not in self.models:
+            return "❌ No model loaded. Please select a model first."
 
-    def resume_training_from_7k_to_8k():
-        """Resume training from 7k model to create 8k model."""
         try:
-            if not MODULES_AVAILABLE:
-                return "❌ Required modules not available. Please check deployment."
-
-            # Import required modules
-            import json
-            import time
-            from datetime import datetime
-
-            import torch
-            from huggingface_hub import HfApi, create_repo, snapshot_download, whoami
-            from model import GPTConfig, GPTModel
-            from train_model import TextDataLoader
-
-            print("🚀 Resuming Training from 7k to 8k Model")
-            print("=" * 50)
-
-            # Configuration
-            hf_model_id = "lemms/openllm-small-extended-7k"
-            additional_steps = 1000  # Train for 1000 more steps to reach 8k
-            total_steps = 8000  # Total steps for the new model
-
-            print(f"📥 Source Model: {hf_model_id}")
-            print(f"📈 Additional Steps: {additional_steps}")
-            print(f"🎯 Target Steps: {total_steps}")
-
-            # Setup authentication
-            print("🔐 Setting up authentication...")
-            try:
-                user_info = whoami()
-                username = user_info.get("name", "unknown")
-                print(f"✅ Authentication successful! User: {username}")
-            except Exception as e:
-                return f"❌ Authentication failed: {e}"
-
-            # Setup device
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            print(f"🖥️ Using device: {device}")
-
-            # Load model from Hugging Face
-            print(f"📥 Loading model from Hugging Face: {hf_model_id}")
-            try:
-                local_dir = snapshot_download(
-                    repo_id=hf_model_id,
-                    repo_type="model",
-                    local_dir=f"downloaded_models/{hf_model_id.replace('/', '_')}",
-                )
-                print(f"✅ Model downloaded to: {local_dir}")
-
-                # Load config
-                config_path = Path(local_dir) / "config.json"
-                if config_path.exists():
-                    with open(config_path, "r") as f:
-                        config_data = json.load(f)
-
-                    config = GPTConfig(
-                        vocab_size=config_data.get("vocab_size", 32000),
-                        block_size=config_data.get("block_size", 1024),
-                        n_layer=config_data.get("n_layer", 6),
-                        n_head=config_data.get("n_head", 6),
-                        n_embd=config_data.get("n_embd", 384),
-                    )
-                    print(f"📊 Loaded model config: {config}")
-                else:
-                    config = GPTConfig.small()
-                    config.vocab_size = 32000
-                    print(f"⚠️ Config file not found, using default config")
-
-                # Create model and load weights
-                model = GPTModel(config)
-                model_path = Path(local_dir) / "pytorch_model.bin"
-
-                if model_path.exists():
-                    state_dict = torch.load(model_path, map_location=device)
-                    model.load_state_dict(state_dict)
-                    print(f"✅ Model weights loaded successfully")
-                else:
-                    raise FileNotFoundError(f"Model file not found: {model_path}")
-
-                model = model.to(device)
+            model = self.models[self.current_model]
+            tokenizer = self.tokenizers[self.current_model]
 
-            except Exception as e:
-                return f"❌ Failed to load model from Hugging Face: {e}"
+            # Tokenize input
+            input_ids = tokenizer.encode(prompt)
+            input_tensor = torch.tensor([input_ids], dtype=torch.long)
 
-            # Create data loaders
-            print(f"📊 Loading training data...")
-            tokenizer_path = "data/tokenizer/tokenizer.model"
-
-            train_loader = TextDataLoader(
-                data_file="data/clean/training_data.txt",
-                tokenizer_path=tokenizer_path,
-                seq_len=1024,
-                batch_size=16,
-                shuffle=True,
+            logger.info(f"🎯 Generating text with prompt: '{prompt[:50]}...'")
+            logger.info(
+                f"📊 Parameters: max_length={max_length}, temperature={temperature}, top_k={top_k}, top_p={top_p}"
             )
 
-            print(f"✅ Data loader created")
-
-            # Setup optimizer
-            optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
-
-            # Training loop
-            print(f"\n🔄 Starting training loop...")
-            start_time = time.time()
-            training_history = []
-            best_loss = float("inf")
-
-            try:
-                train_iterator = iter(train_loader)
-
-                for step in range(additional_steps):
-                    # Get batch
-                    try:
-                        batch = next(train_iterator)
-                    except StopIteration:
-                        # Restart data loader if exhausted
-                        train_loader = TextDataLoader(
-                            data_file="data/clean/training_data.txt",
-                            tokenizer_path=tokenizer_path,
-                            seq_len=1024,
-                            batch_size=16,
-                            shuffle=True,
-                        )
-                        train_iterator = iter(train_loader)
-                        batch = next(train_iterator)
-
-                    # Prepare inputs
-                    if isinstance(batch, (list, tuple)):
-                        inputs = batch[0].to(device)
-                        targets = batch[1].to(device) if len(batch) > 1 else None
-                    else:
-                        inputs = batch.to(device)
-                        targets = None
-
-                    # Forward pass
-                    logits, loss = model(inputs, targets)
-
-                    # Backward pass
-                    loss.backward()
-
-                    # Optimizer step
-                    optimizer.step()
-                    optimizer.zero_grad()
-
-                    # Record training history
-                    training_history.append(
-                        {
-                            "step": 7000 + step + 1,  # Continue from step 7000
-                            "loss": loss.item(),
-                            "timestamp": datetime.now().isoformat(),
-                        }
-                    )
-
-                    # Progress reporting
-                    if (step + 1) % 10 == 0:
-                        elapsed = time.time() - start_time
-                        steps_per_sec = (step + 1) / elapsed
-                        eta = (additional_steps - step - 1) / steps_per_sec
-
-                        print(
-                            f"Step {7000 + step + 1}/{total_steps} | "
-                            f"Loss: {loss.item():.4f} | "
-                            f"Speed: {steps_per_sec:.1f} steps/s | "
-                            f"ETA: {eta/60:.1f} min"
-                        )
-
-                    # Evaluation
-                    if (step + 1) % 250 == 0:
-                        model.eval()
-                        total_loss = 0.0
-                        num_batches = 0
-
-                        with torch.no_grad():
-                            for val_batch in train_loader:  # Use same loader for simplicity
-                                if isinstance(val_batch, (list, tuple)):
-                                    val_inputs = val_batch[0].to(device)
-                                    val_targets = (
-                                        val_batch[1].to(device) if len(val_batch) > 1 else None
-                                    )
-                                else:
-                                    val_inputs = val_batch.to(device)
-                                    val_targets = None
-
-                                val_logits, val_loss = model(val_inputs, val_targets)
-                                total_loss += val_loss.item()
-                                num_batches += 1
-
-                                if num_batches >= 5:  # Limit evaluation
-                                    break
-
-                        avg_val_loss = total_loss / num_batches
-                        model.train()
-                        print(f"📊 Validation Loss: {avg_val_loss:.4f}")
-
-                        # Check for best model
-                        if avg_val_loss < best_loss:
-                            best_loss = avg_val_loss
-                            print(f"🏆 New best validation loss: {best_loss:.4f}")
-
-                print(f"\n🎉 Training completed successfully!")
-                print(f"📊 Final Results:")
-                print(f"   - Additional Steps: {additional_steps}")
-                print(f"   - Total Steps: {total_steps}")
-                print(f"   - Final Loss: {loss.item():.4f}")
-                print(f"   - Best Validation Loss: {best_loss:.4f}")
-                print(f"   - Training Time: {(time.time() - start_time)/3600:.2f} hours")
-
-                # Upload model
-                print(f"\n📤 Uploading model to Hugging Face Hub...")
-
-                # Create model directory
-                model_path = Path("./trained_model")
-                model_path.mkdir(exist_ok=True)
-
-                # Save model files
-                torch.save(model.state_dict(), model_path / "pytorch_model.bin")
-
-                # Save config
-                config_dict = {
-                    "model_type": "openllm",
-                    "model_size": "small",
-                    "vocab_size": 32000,
-                    "block_size": 1024,
-                    "n_layer": 6,
-                    "n_head": 6,
-                    "n_embd": 384,
-                    "training_config": {
-                        "model_size": "small",
-                        "training_steps": total_steps,
-                        "additional_steps": additional_steps,
-                        "base_model": hf_model_id,
-                    },
-                    "training_history": training_history,
-                }
-
-                with open(model_path / "config.json", "w") as f:
-                    json.dump(config_dict, f, indent=2)
-
-                # Create model card
-                readme_content = f"""# OpenLLM Small Model - Extended to 8k Steps
-
-This is an OpenLLM small model trained for {total_steps} steps by resuming training from [lemms/openllm-small-extended-7k](https://huggingface.co/lemms/openllm-small-extended-7k).
-
-## Model Details
-
-- **Model Type**: OpenLLM
-- **Size**: small
-- **Training Steps**: {total_steps}
-- **Additional Steps**: {additional_steps}
-- **Base Model**: [lemms/openllm-small-extended-7k](https://huggingface.co/lemms/openllm-small-extended-7k)
-- **Final Loss**: {training_history[-1]['loss']:.4f} if training_history else 'N/A'
-- **Framework**: PyTorch
-- **License**: GPL-3.0
-
-## Training Configuration
-
-```json
-{json.dumps(config_dict, indent=2)}
-```
-
-## Training History
-
-The model was trained with the following key metrics:
-- Best validation loss: {best_loss:.4f}
-- Total training time: {len(training_history)} steps
-- Device used: {device}
-
-## Usage
-
-This model can be used for text generation and language modeling tasks.
-
-## Author
-
-Louis Chua Bean Chong
-
-## License
-
-GPL-3.0
-"""
+            # Generate text
+            with torch.no_grad():
+                output_ids = model.generate(
+                    input_tensor,
+                    max_new_tokens=max_length,
+                    temperature=temperature,
+                    top_k=top_k,
+                    top_p=top_p,
+                    do_sample=True,
+                )
 
-                with open(model_path / "README.md", "w") as f:
-                    f.write(readme_content)
+            # Decode output
+            generated_text = tokenizer.decode(output_ids[0].tolist())
 
-                # Upload to Hugging Face
-                repo_name = "openllm-small-extended-8k"
-                repo_id = f"{username}/{repo_name}"
+            # Remove the input prompt from the output
+            if generated_text.startswith(prompt):
+                generated_text = generated_text[len(prompt) :].strip()
 
-                try:
-                    # Create repository
-                    create_repo(repo_id=repo_id, repo_type="model", exist_ok=True, private=False)
+            logger.info(f"✅ Generated text: '{generated_text[:100]}...'")
+            return generated_text
 
-                    # Upload files
-                    api = HfApi()
-                    api.upload_folder(
-                        folder_path=str(model_path),
-                        repo_id=repo_id,
-                        repo_type="model",
-                        commit_message=f"Add OpenLLM small model extended to {total_steps} steps",
-                    )
+        except Exception as e:
+            error_msg = f"❌ Generation failed: {str(e)}"
+            logger.error(error_msg)
+            import traceback
 
-                    print(f"✅ Model uploaded successfully!")
-                    print(f"🔗 Model URL: https://huggingface.co/{repo_id}")
+            logger.error(f"📋 Full traceback: {traceback.format_exc()}")
+            return error_msg
 
-                    result = f"✅ 8k Model Training completed successfully!\n\n"
-                    result += f"📊 Results:\n"
-                    result += f"   - Base Model: {hf_model_id}\n"
-                    result += f"   - Additional Steps: {additional_steps}\n"
-                    result += f"   - Total Steps: {total_steps}\n"
-                    result += f"   - Final Loss: {loss.item():.4f}\n"
-                    result += f"   - Best Validation Loss: {best_loss:.4f}\n"
-                    result += f"   - Model URL: https://huggingface.co/{repo_id}\n\n"
-                    result += f"🎉 Extended model available at: https://huggingface.co/{repo_id}"
 
-                    return result
+# Initialize the real inference engine
+inference_engine = RealOpenLLMInference()
 
-                except Exception as e:
-                    print(f"❌ Model upload failed: {e}")
-                    return f"⚠️ Training completed but upload failed: {e}"
 
-            except KeyboardInterrupt:
-                print(f"\n⚠️ Training interrupted by user")
-                return "⚠️ Training was interrupted by user"
+def load_model_info(model_id: str) -> str:
+    """Get information about a specific model"""
+    config = inference_engine.model_configs.get(model_id)
+    if config:
+        return f"**{config['name']}**\n\n{config['description']}\n\n**Parameters:** {config['parameters']}\n**Training Steps:** {config['training_steps']:,}"
+    return "❌ Model not found"
 
-        except Exception as e:
-            return f"❌ Error resuming training: {e}"
 
-    def check_space_environment():
-        """Check the Space environment and configuration."""
-        try:
-            # Check if we're in a Space
-            space_vars = ["SPACE_ID", "SPACE_HOST", "SPACE_REPO_ID"]
-            is_space = any(os.getenv(var) for var in space_vars)
+def generate_text_interface(
+    model_id: str, prompt: str, max_length: int, temperature: float, top_k: int, top_p: float
+) -> str:
+    """Gradio interface function for text generation"""
+    try:
+        # Load model if not already loaded
+        if model_id not in inference_engine.models:
+            logger.info(f"🔄 Loading real model: {model_id}")
+            success = inference_engine.load_model_from_hf(model_id)
+            if not success:
+                return f"❌ Failed to load real model: {model_id}"
 
-            # Check HF_TOKEN
-            hf_token = os.getenv("HF_TOKEN")
+        # Generate text
+        result = inference_engine.generate_text(
+            prompt=prompt, max_length=max_length, temperature=temperature, top_k=top_k, top_p=top_p
+        )
 
-            result = "🔍 Space Environment Check:\n\n"
+        return result
 
-            if is_space:
-                result += "✅ Running in Hugging Face Space environment\n"
-                for var in space_vars:
-                    value = os.getenv(var)
-                    if value:
-                        result += f"   - {var}: {value}\n"
-            else:
-                result += "ℹ️ Running in local environment\n"
-
-            # Test Space's built-in authentication
-            try:
-                from huggingface_hub import whoami
-
-                user_info = whoami()
-                result += f"✅ Space built-in authentication working\n"
-                result += f"   - User: {user_info['name']}\n"
-                result += f"   - Full name: {user_info['fullname']}\n"
-                result += f"   - Authentication: Space built-in token\n"
-            except Exception as auth_error:
-                result += f"❌ Space built-in authentication failed: {str(auth_error)[:50]}...\n"
-
-                if hf_token:
-                    result += f"✅ HF access token found: {hf_token[:8]}...{hf_token[-4:]}\n"
-                    result += "   - Source: HF access token in Space settings\n"
-                else:
-                    result += "❌ HF access token not found\n"
-                    result += "   - Please set HF_TOKEN in Space settings with HF access token\n"
-                    result += "   - Or ensure Space has proper authentication permissions\n"
+    except Exception as e:
+        error_msg = f"❌ Error in generation interface: {str(e)}"
+        logger.error(error_msg)
+        return error_msg
 
-            result += f"\n📁 Available modules: {'✅' if MODULES_AVAILABLE else '❌'}"
 
-            return result
+# Create Gradio interface
+def create_interface():
+    """Create the Gradio interface"""
 
-        except Exception as e:
-            return f"❌ Error checking environment: {e}"
-
-    # Create the Gradio interface with security mitigations
-    with gr.Blocks(
-        title="OpenLLM Training Space",
-        theme=gr.themes.Soft(),
-        # Security mitigations
-        analytics_enabled=False,  # Disable analytics
-    ) as interface:
+    with gr.Blocks(title="🚀 OpenLLM Real Models Space", theme=gr.themes.Soft()) as interface:
+        # Header
         gr.Markdown(
             """
-        # 🚀 OpenLLM Training Space
+        # 🚀 OpenLLM Real Models Space
+        
+        Welcome to the OpenLLM Real Models Space! This interface uses **actual trained models** from Hugging Face.
         
-        Welcome to the OpenLLM Training Space! This Space provides a complete environment for training OpenLLM models with automatic Hugging Face authentication and model upload.
+        ## 🎯 Real Trained Models
         
-                 ## 🔐 Authentication
+        We provide **5 different real models** with varying training steps:
         
-         This Space uses HF access token for secure authentication. The HF_TOKEN is automatically available from your Space settings.
+        | Model | Training Steps | Parameters | Performance |
+        |-------|---------------|------------|-------------|
+        | **4k Model** | 4,000 | 35.8M | Early training stage |
+        | **6k Model** | 6,000 | 35.8M | Improved coherence (Perplexity: 816.040) |
+        | **7k Model** | 7,000 | 35.8M | Enhanced quality (Loss: 2.100, Perplexity: 8.200) |
+        | **8k Model** | 8,000 | 35.8M | Sophisticated understanding |
+        | **9k Model** | 9,000 | 35.8M | Best performing model |
+        | **10k Model** | 10,000 | 35.8M | Latest extended training |
         
-        ## 📋 Available Actions
+        **These are real GPT-style transformer models trained on Wikipedia passages from the SQuAD dataset.**
         
-        1. **Environment Check**: Verify Space configuration and authentication
-        2. **Authentication Test**: Test Hugging Face authentication
-        3. **Training Interface**: Unified interface for fresh training and resume training
+        ---
         """
         )
 
-        with gr.Tab("🔍 Environment Check"):
-            gr.Markdown("Check the Space environment and configuration.")
-            env_check_btn = gr.Button("Check Environment", variant="primary")
-            env_output = gr.Textbox(label="Environment Status", lines=10, interactive=False)
-            env_check_btn.click(check_space_environment, outputs=env_output)
-
-        with gr.Tab("🔐 Authentication Test"):
-            gr.Markdown("Test Hugging Face authentication using HF access token.")
-            auth_test_btn = gr.Button("Run Authentication Test", variant="primary")
-            auth_output = gr.Textbox(label="Authentication Results", lines=15, interactive=False)
-            auth_test_btn.click(run_authentication_test, outputs=auth_output)
-
-        with gr.Tab("🚀 Training Interface"):
-            gr.Markdown(
-                """
-            # 🚀 OpenLLM Training Interface
-            
-            Choose your training mode and configure parameters for model training.
-            
-            ## 🎯 Training Modes
-            
-            **1. Fresh Training**: Start training from scratch with a new model
-            **2. Resume Training**: Load the 7k model and continue training to 8k steps
-            
-            ## 📊 Training Parameters
-            
-            - **Model Size**: Choose the model size (small, medium, large)
-            - **Training Steps**: Number of training steps (default: 8000)
-            - **Training Mode**: Select between fresh training or resume training
-            - **Real Training**: Enable comprehensive training with checkpoints and validation
-            
-            ## 🎉 Expected Results
-            
-            - Training will complete successfully
-            - Model will be uploaded to Hugging Face Hub
-            - Repository will be created with proper model files
-            """
-            )
-
-            with gr.Row():
-                training_mode = gr.Radio(
-                    choices=["Fresh Training", "Resume 7k to 8k"],
-                    value="Fresh Training",
-                    label="Training Mode",
-                    info="Choose between fresh training or resuming from 7k model",
+        with gr.Row():
+            with gr.Column(scale=1):
+                # Model selection
+                model_dropdown = gr.Dropdown(
+                    choices=list(inference_engine.model_configs.keys()),
+                    value="openllm-small-extended-10k",
+                    label="🎯 Select Model",
+                    info="Choose the real trained model to use",
                 )
 
-            with gr.Row():
-                model_size = gr.Dropdown(
-                    choices=["small", "medium", "large"],
-                    value="small",
-                    label="Model Size",
-                    info="Choose the model size for training (only applies to fresh training)",
-                    interactive=True,
+                # Model information display
+                model_info = gr.Markdown(
+                    value=load_model_info("openllm-small-extended-10k"), label="📋 Model Information"
                 )
-                training_steps = gr.Number(
-                    value=8000,
-                    label="Training Steps",
-                    info="Number of training steps (only applies to fresh training)",
-                    minimum=1000,
-                    maximum=50000,
-                    interactive=True,
+
+                # Update model info when selection changes
+                model_dropdown.change(
+                    fn=load_model_info, inputs=[model_dropdown], outputs=[model_info]
                 )
 
-            with gr.Row():
-                use_real_training = gr.Checkbox(
-                    value=False,
-                    label="Use Real Training",
-                    info="Enable real model training with checkpoints and validation (slower but more realistic)",
+            with gr.Column(scale=2):
+                # Input prompt
+                prompt_input = gr.Textbox(
+                    lines=5,
+                    label="📝 Input Prompt",
+                    placeholder="Enter your text prompt here...",
+                    info="The text that will be used as input for generation",
                 )
 
-            # Function to update UI based on training mode
-            def update_ui_for_mode(mode):
-                """Update UI elements based on selected training mode."""
-                if mode == "Resume 7k to 8k":
-                    return gr.Dropdown(interactive=False), gr.Number(interactive=False)
-                else:
-                    return gr.Dropdown(interactive=True), gr.Number(interactive=True)
+                # Generation parameters
+                with gr.Row():
+                    max_length = gr.Slider(
+                        minimum=10,
+                        maximum=500,
+                        value=100,
+                        step=10,
+                        label="📏 Max Length",
+                        info="Maximum number of tokens to generate",
+                    )
 
-            # Update UI when training mode changes
-            training_mode.change(
-                update_ui_for_mode,
-                inputs=[training_mode],
-                outputs=[model_size, training_steps],
-            )
+                    temperature = gr.Slider(
+                        minimum=0.1,
+                        maximum=2.0,
+                        value=0.7,
+                        step=0.1,
+                        label="🌡️ Temperature",
+                        info="Controls randomness (higher = more random)",
+                    )
 
-            # Unified training function
-            def unified_training(mode, model_size, training_steps, use_real_training):
-                """Unified training function that handles both fresh and resume training."""
-                if mode == "Resume 7k to 8k":
-                    return resume_training_from_7k_to_8k()
-                else:
-                    return run_training(model_size, training_steps, use_real_training)
+                with gr.Row():
+                    top_k = gr.Slider(
+                        minimum=1,
+                        maximum=100,
+                        value=50,
+                        step=1,
+                        label="🔝 Top-K",
+                        info="Number of highest probability tokens to consider",
+                    )
 
-            train_btn = gr.Button("Start Training", variant="primary", size="lg")
-            train_output = gr.Textbox(label="Training Results", lines=25, interactive=False)
+                    top_p = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.9,
+                        step=0.1,
+                        label="📊 Top-P",
+                        info="Nucleus sampling parameter",
+                    )
 
-            train_btn.click(
-                unified_training,
-                inputs=[training_mode, model_size, training_steps, use_real_training],
-                outputs=train_output,
-            )
+                # Generate button
+                generate_btn = gr.Button("🚀 Generate Text", variant="primary", size="lg")
+
+        # Output
+        output_text = gr.Textbox(
+            lines=10, label="🎯 Generated Text", info="The generated text will appear here"
+        )
+
+        # Connect the generate button
+        generate_btn.click(
+            fn=generate_text_interface,
+            inputs=[model_dropdown, prompt_input, max_length, temperature, top_k, top_p],
+            outputs=[output_text],
+        )
 
-        with gr.Tab("📚 Documentation"):
-            gr.Markdown(
-                """
-            ## 📖 Available Documentation
-            
-            - **HUGGINGFACE_SPACE_SETUP_GUIDE.md**: Complete setup guide
-            - **SPACE_AUTHENTICATION_SUMMARY.md**: Authentication summary
-            - **SPACE_READY_SUMMARY.md**: Deployment summary
-            
-            ## 🔧 Available Scripts
-            
-            - **space_auth_test.py**: Authentication verification
-            - **openllm_training_with_auth.py**: Complete training script
-            - **integrate_auth_into_training.py**: Integration guide
-            - **setup_hf_space_auth.py**: Space authentication setup
-            - **verify_space_auth.py**: Space verification script
-            
-            ## 🎯 Quick Start
-            
-            1. Check the environment to verify configuration
-            2. Run authentication test to ensure GitHub secrets are working
-            3. Start training with your desired parameters
-            4. Monitor the training progress and model upload
-            
-            ## 🔒 Security
-            
-            - HF_TOKEN is securely stored in GitHub repository secrets
-            - No hardcoded tokens in any scripts
-            - Automatic cleanup of test repositories
-            - Proper error handling and logging
+        # Footer
+        gr.Markdown(
             """
-            )
+        ---
+        
+        ## 🔧 Technical Details
+        
+        - **Architecture**: GPT-style transformer decoder
+        - **Model Size**: Small (6 layers, 8 heads, 512 embedding dim)
+        - **Vocabulary**: 32k tokens (SentencePiece BPE)
+        - **Training Data**: Wikipedia passages from SQuAD dataset
+        - **Framework**: PyTorch with real trained models
+        - **Gradio Version**: 4.44.1 (latest)
+        
+        **These models generate actual text based on their training on Wikipedia content.**
+        
+        **Model Sources:**
+        - [4k Model](https://huggingface.co/lemms/openllm-small-extended-4k)
+        - [6k Model](https://huggingface.co/lemms/openllm-small-extended-6k)
+        - [7k Model](https://huggingface.co/lemms/openllm-small-extended-7k)
+        - [8k Model](https://huggingface.co/lemms/openllm-small-extended-8k)
+        - [9k Model](https://huggingface.co/lemms/openllm-small-extended-9k)
+        - [10k Model](https://huggingface.co/lemms/openllm-small-extended-10k)
+        """
+        )
 
     return interface
 
 
+# Create and launch the interface
 if __name__ == "__main__":
-    # Create and launch the interface
-    interface = create_space_interface()
-    interface.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        # Security mitigations for Gradio vulnerabilities
-        allowed_paths=[],  # Restrict file access
-        auth=None,  # Disable authentication to prevent code injection
-        quiet=True,  # Reduce logging
-    )
+    interface = create_interface()
+    interface.launch(server_name="0.0.0.0", server_port=7860, share=False, debug=True)