app.py

import gradio as gr
import torch
import torch.nn as nn
from model import SmolLM2_135M  # Import your model class
import yaml

# Device setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load model
@torch.no_grad()
def load_model():
    """Load the trained model"""
    print("Loading model...")
    
    # Load config
    with open('config.yaml', 'r') as f:
        config = yaml.safe_load(f)
    
    # Initialize model
    model = SmolLM2_135M(
        vocab_size=config['vocab_size'],
        d_model=config['d_model'],
        n_layers=config['n_layers'],
        n_heads=config['n_heads'],
        # Add other config parameters
    ).to(device)
    
    # Load checkpoint
    checkpoint = torch.load('checkpoints/checkpoint_step_5050.pt', 
                           map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    
    print(f"Model loaded successfully on {device}")
    return model, checkpoint

# Load model at startup
model, checkpoint = load_model()

# Tokenizer (adjust based on your implementation)
def tokenize(text, max_length=128):
    """Simple character-level tokenizer - REPLACE with your actual tokenizer"""
    # This is a placeholder - use your actual tokenizer
    tokens = [ord(c) for c in text[:max_length]]
    return torch.tensor(tokens).unsqueeze(0).to(device)

def detokenize(tokens):
    """Convert tokens back to text - REPLACE with your actual detokenizer"""
    # This is a placeholder - use your actual detokenizer
    return ''.join([chr(t) for t in tokens if t < 128])

@torch.no_grad()
def generate_text(
    prompt,
    max_length=100,
    temperature=0.8,
    top_k=50,
    top_p=0.9
):
    """Generate text from prompt"""
    try:
        # Tokenize input
        input_ids = tokenize(prompt)
        
        # Generate
        generated = input_ids[0].tolist()
        
        for _ in range(max_length):
            # Get model predictions
            input_tensor = torch.tensor([generated]).to(device)
            logits = model(input_tensor)
            
            # Get next token logits
            next_token_logits = logits[0, -1, :] / temperature
            
            # Apply top-k filtering
            if top_k > 0:
                indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
                next_token_logits[indices_to_remove] = float('-inf')
            
            # Apply top-p (nucleus) filtering
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
                cumulative_probs = torch.cumsum(torch.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[sorted_indices_to_remove]
                next_token_logits[indices_to_remove] = float('-inf')
            
            # Sample next token
            probs = torch.softmax(next_token_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1).item()
            
            generated.append(next_token)
            
            # Stop if EOS token (adjust based on your vocab)
            if next_token == 0:  # Assuming 0 is EOS
                break
        
        # Detokenize
        output_text = detokenize(generated)
        return output_text
        
    except Exception as e:
        return f"Error generating text: {str(e)}"

def get_model_info():
    """Display model information"""
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    info = f"""
    ### 📊 Model Information
    
    **Total Parameters:** {total_params:,} (~{total_params/1e6:.1f}M)
    **Trainable Parameters:** {trainable_params:,}
    **Training Steps:** {checkpoint.get('step', 'N/A')}
    **Device:** {device}
    **Model Architecture:** SmolLM2-135M
    
    ### 🎯 Training Details
    - Trained for 5,000 steps
    - Checkpoint saved and reloaded
    - Additional 50 steps after reload
    - Predictions logged every 500 steps
    """
    return info

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🤖 SmolLM2-135M: From-Scratch Implementation
    
    This is a complete reverse-engineered implementation of SmolLM2-135M, trained from scratch.
    
    **GitHub:** [abi2024/smollm2-135-implementation](https://github.com/abi2024/smollm2-135-implementation)
    """)
    
    with gr.Tab("🎮 Generate Text"):
        with gr.Row():
            with gr.Column():
                prompt_input = gr.Textbox(
                    label="Prompt",
                    placeholder="Enter your prompt here...",
                    lines=3,
                    value="Once upon a time"
                )
                
                with gr.Row():
                    max_length_slider = gr.Slider(
                        minimum=10,
                        maximum=500,
                        value=100,
                        step=10,
                        label="Max Length"
                    )
                    temperature_slider = gr.Slider(
                        minimum=0.1,
                        maximum=2.0,
                        value=0.8,
                        step=0.1,
                        label="Temperature"
                    )
                
                with gr.Row():
                    top_k_slider = gr.Slider(
                        minimum=0,
                        maximum=100,
                        value=50,
                        step=5,
                        label="Top-K"
                    )
                    top_p_slider = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=0.9,
                        step=0.05,
                        label="Top-P"
                    )
                
                generate_btn = gr.Button("🚀 Generate", variant="primary")
            
            with gr.Column():
                output_text = gr.Textbox(
                    label="Generated Text",
                    lines=10,
                    interactive=False
                )
        
        generate_btn.click(
            fn=generate_text,
            inputs=[
                prompt_input,
                max_length_slider,
                temperature_slider,
                top_k_slider,
                top_p_slider
            ],
            outputs=output_text
        )
        
        gr.Markdown("""
        ### 💡 Tips:
        - **Temperature**: Higher = more creative, Lower = more focused
        - **Top-K**: Limits vocabulary to K most likely tokens
        - **Top-P**: Nucleus sampling - cumulative probability threshold
        """)
    
    with gr.Tab("📊 Model Info"):
        model_info_display = gr.Markdown(get_model_info())
        
        gr.Markdown("""
        ### 🏗️ Architecture Details
        
        This model was reverse-engineered by:
        1. Analyzing the official SmolLM2 repository
        2. Extracting architecture from pretrained weights
        3. Implementing from scratch in PyTorch
        4. Validating by swapping weights with pretrained model
        
        ### ⚡ Optimizations Used
        - Flash Attention 2
        - Mixed Precision Training (BF16/FP16)
        - Gradient Accumulation
        - torch.compile()
        
        ### 📈 Training Process
        - **Step 0-5000**: Main training with periodic predictions
        - **Checkpoint**: Model saved and reloaded to validate state preservation
        - **Step 5000-5050**: Continued training to test checkpoint robustness
        """)
    
    with gr.Tab("🎯 Example Prompts"):
        gr.Markdown("""
        ### Try these prompts:
        
        1. **Story Generation**
```
           Once upon a time in a land far away
```
        
        2. **Code Completion**
```
           def fibonacci(n):
```
        
        3. **Question Answering**
```
           Q: What is machine learning?
           A:
```
        
        4. **Creative Writing**
```
           The old house at the end of the street was
```
        
        5. **Technical Explanation**
```
           Neural networks work by
```
        """)

# Launch
if __name__ == "__main__":
    demo.launch()