Upload app.py with huggingface_hub

app.py

@@ -1,426 +1,627 @@
-#!/usr/bin/env python3
-"""
-OpenLLM Inference Space - Gradio Interface
-Loads models from Hugging Face repositories to avoid storage limits
-"""
-
-import gradio as gr
-import torch
-import json
-import os
-from pathlib import Path
-from typing import Dict, Any, Optional
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Import our custom modules
-try:
-    from core.src.model import GPTConfig, GPTModel
-    from core.src.inference_server import OpenLLMInference
-    logger.info("✅ Successfully imported core modules")
-except ImportError as e:
-    logger.error(f"❌ Failed to import core modules: {e}")
-    raise
-
-class OpenLLMInferenceEngine:
-    """
-    Inference engine that loads models from Hugging Face repositories
-    """
-    
-    def __init__(self):
-        self.models = {}
-        self.tokenizers = {}
-        self.current_model = None
-        self.current_tokenizer = None
-        
-        # Model configurations with Hugging Face repository IDs
-        self.model_configs = {
-            "openllm-small-extended-4k": {
-                "name": "OpenLLM Small (4k steps)",
-                "description": "Small model trained for 4,000 steps - Early training stage",
-                "hf_repo": "lemms/openllm-small-extended-4k",
-                "local_path": "models/small-extended-4k",  # Fallback if HF repo fails
-                "checkpoint": "best_model.pt",
-                "config": "config.json"
-            },
-            "openllm-small-extended-6k": {
-                "name": "OpenLLM Small (6k steps)", 
-                "description": "Small model trained for 6,000 steps - Improved coherence",
-                "hf_repo": "lemms/openllm-small-extended-6k",
-                "local_path": "models/small-extended-6k",
-                "checkpoint": "best_model.pt",
-                "config": "config.json"
-            },
-            "openllm-small-extended-7k": {
-                "name": "OpenLLM Small (7k steps)",
-                "description": "Small model trained for 7,000 steps - Enhanced quality",
-                "hf_repo": "lemms/openllm-small-extended-7k", 
-                "local_path": "models/small-extended-7k",
-                "checkpoint": "best_model.pt",
-                "config": "config.json"
-            },
-            "openllm-small-extended-8k": {
-                "name": "OpenLLM Small (8k steps)",
-                "description": "Small model trained for 8,000 steps - Sophisticated understanding",
-                "hf_repo": "lemms/openllm-small-extended-8k",
-                "local_path": "models/small-extended-8k", 
-                "checkpoint": "best_model.pt",
-                "config": "config.json"
-            },
-            "openllm-small-extended-9k": {
-                "name": "OpenLLM Small (9k steps)",
-                "description": "Small model trained for 9,000 steps - Best performing model",
-                "hf_repo": "lemms/openllm-small-extended-9k",
-                "local_path": "models/small-extended-9k",
-                "checkpoint": "best_model.pt", 
-                "config": "config.json"
-            },
-            "openllm-small-extended-10k": {
-                "name": "OpenLLM Small (10k steps)",
-                "description": "Small model trained for 10,000 steps - Latest extended training",
-                "hf_repo": "lemms/openllm-small-extended-10k",
-                "local_path": "models/small-extended-10k",
-                "checkpoint": "best_model.pt", 
-                "config": "config.json"
-            }
-        }
-        
-        logger.info("🚀 OpenLLM Inference Engine initialized")
-        logger.info(f"📋 Available models: {list(self.model_configs.keys())}")
-    
-    def load_model_from_hf(self, model_id: str) -> bool:
-        """
-        Load model from Hugging Face repository
-        """
-        try:
-            from huggingface_hub import snapshot_download
-            
-            config = self.model_configs.get(model_id)
-            if not config:
-                logger.error(f"❌ Unknown model ID: {model_id}")
-                return False
-            
-            logger.info(f"📥 Loading model from HF: {config['hf_repo']}")
-            
-            # Download model files 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"]
-            )
-            
-            logger.info(f"✅ Downloaded model to: {local_dir}")
-            
-            # Initialize inference engine with downloaded model
-            inference_engine = OpenLLMInference(
-                model_path=Path(local_dir),
-                model_format="pytorch"
-            )
-            
-            self.models[model_id] = inference_engine
-            self.current_model = model_id
-            
-            logger.info(f"✅ Successfully loaded model: {model_id}")
-            return True
-            
-        except Exception as e:
-            logger.error(f"❌ Failed to load model from HF {model_id}: {e}")
-            
-            # Fallback to local model if available
-            logger.info(f"🔄 Trying fallback to local model...")
-            return self.load_model_local(model_id)
-    
-    def load_model_local(self, model_id: str) -> bool:
-        """
-        Load model from local files (fallback method)
-        """
-        try:
-            config = self.model_configs.get(model_id)
-            if not config:
-                return False
-            
-            local_path = Path(config['local_path'])
-            if not local_path.exists():
-                logger.error(f"❌ Local model path not found: {local_path}")
-                return False
-            
-            logger.info(f"📂 Loading local model: {local_path}")
-            
-            # Initialize inference engine with local model
-            inference_engine = OpenLLMInference(
-                model_path=local_path,
-                model_format="pytorch"
-            )
-            
-            self.models[model_id] = inference_engine
-            self.current_model = model_id
-            
-            logger.info(f"✅ Successfully loaded local model: {model_id}")
-            return True
-            
-        except Exception as e:
-            logger.error(f"❌ Failed to load local model {model_id}: {e}")
-            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 currently loaded model
-        """
-        if not self.current_model or self.current_model not in self.models:
-            return "❌ No model loaded. Please select a model first."
-        
-        try:
-            logger.info(f"🎯 Generating text with model: {self.current_model}")
-            logger.info(f"📝 Prompt: {prompt[:100]}...")
-            logger.info(f"⚙️  Parameters: max_length={max_length}, temp={temperature}, top_k={top_k}, top_p={top_p}")
-            
-            # Get the inference engine
-            inference_engine = self.models[self.current_model]
-            
-            # Prepare generation parameters
-            generation_params = {
-                "max_length": max_length,
-                "temperature": temperature,
-                "top_k": top_k,
-                "top_p": top_p,
-                "do_sample": True,
-                "pad_token_id": 0,
-                "eos_token_id": 1
-            }
-            
-            # Generate text
-            generated_texts = inference_engine.generate(
-                prompt=prompt,
-                **generation_params
-            )
-            
-            # Extract the generated text
-            if isinstance(generated_texts, list) and len(generated_texts) > 0:
-                generated_text = generated_texts[0]
-            else:
-                generated_text = str(generated_texts)
-            
-            logger.info(f"✅ Generated text length: {len(generated_text)}")
-            return generated_text
-            
-        except Exception as e:
-            error_msg = f"❌ Generation failed: {str(e)}"
-            logger.error(error_msg)
-            return error_msg
-
-# Initialize the inference engine
-inference_engine = OpenLLMInferenceEngine()
-
-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']}"
-    return "❌ Model not found"
-
-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 model: {model_id}")
-            success = inference_engine.load_model_from_hf(model_id)
-            if not success:
-                return f"❌ Failed to load model: {model_id}"
-        
-        # Generate text
-        result = inference_engine.generate_text(
-            prompt=prompt,
-            max_length=max_length,
-            temperature=temperature,
-            top_k=top_k,
-            top_p=top_p
-        )
-        
-        return result
-        
-    except Exception as e:
-        error_msg = f"❌ Error in generation interface: {str(e)}"
-        logger.error(error_msg)
-        return error_msg
-
-# Create Gradio interface
-def create_interface():
-    """Create the Gradio interface"""
-    
-    with gr.Blocks(
-        title="🚀 OpenLLM Inference Space",
-        theme=gr.themes.Soft(),
-        css="""
-        .gradio-container {
-            max-width: 1200px !important;
-            margin: auto !important;
-        }
-        .model-info {
-            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-            color: white;
-            padding: 20px;
-            border-radius: 10px;
-            margin-bottom: 20px;
-        }
-        """
-    ) as interface:
-        
-        # Header
-        gr.Markdown("""
-        # 🚀 OpenLLM Inference Space
-        
-        Welcome to the OpenLLM Inference Space! This is a comprehensive interface for running inference on our trained OpenLLM models with customizable parameters.
-        
-        ## 🎯 Available Models
-        
-        We provide **5 different models** trained for varying numbers of steps:
-        
-        | Model | Training Steps | Description | Best Loss |
-        |-------|---------------|-------------|-----------|
-        | **4k Model** | 4,000 | Early training stage, basic language patterns | ~6.2 |
-        | **6k Model** | 6,000 | Improved coherence, better vocabulary usage | ~5.8 |
-        | **7k Model** | 7,000 | Enhanced text generation quality | ~5.5 |
-        | **8k Model** | 8,000 | More sophisticated language understanding | ~5.3 |
-        | **9k Model** | 9,000 | Best performing model (latest training) | ~5.2 |
-        
-        ---
-        """)
-        
-        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-9k",
-                    label="🎯 Select Model",
-                    info="Choose the model to use for inference"
-                )
-                
-                # Model information display
-                model_info = gr.Markdown(
-                    value=load_model_info("openllm-small-extended-9k"),
-                    label="📋 Model Information"
-                )
-                
-                # Update model info when selection changes
-                model_dropdown.change(
-                    fn=load_model_info,
-                    inputs=[model_dropdown],
-                    outputs=[model_info]
-                )
-            
-            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"
-                )
-                
-                # 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"
-                    )
-                    
-                    temperature = gr.Slider(
-                        minimum=0.1,
-                        maximum=2.0,
-                        value=0.7,
-                        step=0.1,
-                        label="🌡️ Temperature",
-                        info="Controls randomness (higher = more random)"
-                    )
-                
-                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"
-                    )
-                    
-                    top_p = gr.Slider(
-                        minimum=0.1,
-                        maximum=1.0,
-                        value=0.9,
-                        step=0.1,
-                        label="📊 Top-P",
-                        info="Nucleus sampling parameter"
-                    )
-                
-                # 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]
-        )
-        
-        # 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**: Custom dataset with 27.6M tokens
-        - **Framework**: PyTorch with custom training loop
-        
-        ## 📚 Usage Tips
-        
-        1. **Start with the 9k model** for best results
-        2. **Lower temperature** (0.3-0.7) for more focused output
-        3. **Higher temperature** (0.8-1.2) for more creative output
-        4. **Adjust max_length** based on your needs
-        5. **Experiment with top-k and top-p** for different generation styles
-        
-        ---
-        
-        **Built with ❤️ using OpenLLM Framework**
-        """)
-    
-    return interface
-
-# Create and launch the interface
-if __name__ == "__main__":
-    interface = create_interface()
-    interface.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        debug=True
-    )
+#!/usr/bin/env python3
+"""
+OpenLLM Real Models App - Final working version with correct attribute naming
+"""
+
+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=True, **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),
+        ))
+        
+        # Language model head - MUST have bias to match saved model
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=True)
+        
+        # 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')
+            
+            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)
+        
+        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"
+            }
+        }
+        
+        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:
+            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"]
+            )
+            
+            logger.info(f"✅ Downloaded model to: {local_dir}")
+            
+            # Load model and tokenizer
+            success = self._load_model_and_tokenizer(local_dir, model_id)
+            if success:
+                self.current_model = model_id
+                logger.info(f"✅ Successfully loaded real model: {model_id}")
+                return True
+            else:
+                return False
+                
+        except Exception as e:
+            logger.error(f"❌ Failed to load real model from HF {model_id}: {e}")
+            return False
+    
+    def _load_model_and_tokenizer(self, model_dir: str, model_id: str) -> bool:
+        """Load model and tokenizer from local directory"""
+        try:
+            model_path = Path(model_dir)
+            
+            # 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)
+                
+                logger.info(f"📋 Config data keys: {list(config_data.keys())}")
+                
+                # Handle different config structures
+                if 'model_config' in config_data:
+                    # Extract model_config section
+                    model_config_data = config_data['model_config']
+                else:
+                    # 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'
+                }
+                
+                config_kwargs = {}
+                for key, value in model_config_data.items():
+                    if key in expected_params:
+                        config_kwargs[key] = value
+                
+                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=True
+                )
+            
+            # 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
+            
+            return True
+            
+        except Exception as 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."
+        
+        try:
+            model = self.models[self.current_model]
+            tokenizer = self.tokenizers[self.current_model]
+            
+            # Tokenize input
+            input_ids = tokenizer.encode(prompt)
+            input_tensor = torch.tensor([input_ids], dtype=torch.long)
+            
+            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}")
+            
+            # 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
+                )
+            
+            # Decode output
+            generated_text = tokenizer.decode(output_ids[0].tolist())
+            
+            # Remove the input prompt from the output
+            if generated_text.startswith(prompt):
+                generated_text = generated_text[len(prompt):].strip()
+            
+            logger.info(f"✅ Generated text: '{generated_text[:100]}...'")
+            return generated_text
+            
+        except Exception as e:
+            error_msg = f"❌ Generation failed: {str(e)}"
+            logger.error(error_msg)
+            import traceback
+            logger.error(f"📋 Full traceback: {traceback.format_exc()}")
+            return error_msg
+
+# Initialize the real inference engine
+inference_engine = RealOpenLLMInference()
+
+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"
+
+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}"
+        
+        # Generate text
+        result = inference_engine.generate_text(
+            prompt=prompt,
+            max_length=max_length,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p
+        )
+        
+        return result
+        
+    except Exception as e:
+        error_msg = f"❌ Error in generation interface: {str(e)}"
+        logger.error(error_msg)
+        return error_msg
+
+# Create Gradio interface
+def create_interface():
+    """Create the Gradio interface"""
+    
+    with gr.Blocks(
+        title="🚀 OpenLLM Real Models Space",
+        theme=gr.themes.Soft()
+    ) as interface:
+        
+        # Header
+        gr.Markdown("""
+        # 🚀 OpenLLM Real Models Space
+        
+        Welcome to the OpenLLM Real Models Space! This interface uses **actual trained models** from Hugging Face.
+        
+        ## 🎯 Real Trained Models
+        
+        We provide **5 different real models** with varying training steps:
+        
+        | 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 |
+        
+        **These are real GPT-style transformer models trained on Wikipedia passages from the SQuAD dataset.**
+        
+        ---
+        """)
+        
+        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-9k",
+                    label="🎯 Select Model",
+                    info="Choose the real trained model to use"
+                )
+                
+                # Model information display
+                model_info = gr.Markdown(
+                    value=load_model_info("openllm-small-extended-9k"),
+                    label="📋 Model Information"
+                )
+                
+                # Update model info when selection changes
+                model_dropdown.change(
+                    fn=load_model_info,
+                    inputs=[model_dropdown],
+                    outputs=[model_info]
+                )
+            
+            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"
+                )
+                
+                # 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"
+                    )
+                    
+                    temperature = gr.Slider(
+                        minimum=0.1,
+                        maximum=2.0,
+                        value=0.7,
+                        step=0.1,
+                        label="🌡️ Temperature",
+                        info="Controls randomness (higher = more random)"
+                    )
+                
+                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"
+                    )
+                    
+                    top_p = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.9,
+                        step=0.1,
+                        label="📊 Top-P",
+                        info="Nucleus sampling parameter"
+                    )
+                
+                # 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]
+        )
+        
+        # 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)
+        """)
+    
+    return interface
+
+# Create and launch the interface
+if __name__ == "__main__":
+    interface = create_interface()
+    interface.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        debug=True
+    )
+