Initial Gradio demo upload

README.md

@@ -1,12 +1,52 @@
 ---
-title: Tiny Llm Demo
-emoji: 💻
+title: Tiny-LLM Text Generator
+emoji: 🤖
 colorFrom: blue
-colorTo: red
+colorTo: purple
 sdk: gradio
-sdk_version: 6.5.1
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
+license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Tiny-LLM Text Generator
+
+A **54 million parameter** language model trained **from scratch** on Wikipedia.
+
+## About
+
+This demonstrates that meaningful language models can be trained on consumer hardware with modest compute budgets!
+
+## Architecture
+
+| Component | Value |
+|-----------|-------|
+| Parameters | 54.93M |
+| Layers | 12 |
+| Hidden Size | 512 |
+| Attention Heads | 8 |
+| Intermediate (FFN) | 1408 |
+| Vocab Size | 32,000 |
+| Max Sequence Length | 512 |
+| Position Encoding | RoPE |
+| Normalization | RMSNorm |
+| Activation | SwiGLU |
+
+## Training
+
+- **Training Steps**: 50,000
+- **Tokens**: ~100M
+- **Hardware**: NVIDIA RTX 5090 (32GB)
+- **Training Time**: ~3 hours
+
+## Model
+
+[jonmabe/tiny-llm-54m](https://huggingface.co/jonmabe/tiny-llm-54m)
+
+## Limitations
+
+- Small model size limits knowledge and capabilities
+- Trained only on Wikipedia - limited domain coverage
+- May generate factually incorrect information
+- Not instruction-tuned

app.py

@@ -0,0 +1,254 @@
+"""
+Tiny-LLM Demo - Text Generation with a 54M Parameter Model
+
+This model was trained from scratch on Wikipedia data.
+"""
+
+import gradio as gr
+import torch
+from huggingface_hub import hf_hub_download
+from model import TinyLLM, MODEL_CONFIG
+
+# Model configuration
+MODEL_ID = "jonmabe/tiny-llm-54m"
+MODEL_FILENAME = "final_model.pt"
+
+# Try to use transformers tokenizer, fall back to simple tokenizer
+try:
+    from transformers import AutoTokenizer
+    # Try to load from model repo, fall back to GPT-2 tokenizer
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    except:
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    USE_HF_TOKENIZER = True
+except Exception as e:
+    print(f"Could not load HuggingFace tokenizer: {e}")
+    USE_HF_TOKENIZER = False
+    tokenizer = None
+
+# Load model
+print("Downloading model...")
+model_path = hf_hub_download(repo_id=MODEL_ID, filename=MODEL_FILENAME)
+print(f"Model downloaded to {model_path}")
+
+print("Loading model...")
+checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
+
+# Get config from checkpoint if available
+if "config" in checkpoint and isinstance(checkpoint["config"], dict):
+    config = checkpoint["config"]
+    if "model" in config:
+        config = config["model"]
+else:
+    config = MODEL_CONFIG
+
+# Initialize model
+model = TinyLLM(config)
+
+# Load weights
+if "model_state_dict" in checkpoint:
+    state_dict = checkpoint["model_state_dict"]
+else:
+    state_dict = checkpoint
+
+missing, unexpected = model.load_state_dict(state_dict, strict=False)
+if missing:
+    print(f"Warning: Missing keys: {missing[:5]}...")
+if unexpected:
+    print(f"Warning: Unexpected keys: {unexpected[:5]}...")
+
+# Move to device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+model.eval()
+
+total_params = sum(p.numel() for p in model.parameters())
+print(f"Model loaded on {device} with {total_params:,} parameters")
+
+
+def generate_text(
+    prompt: str,
+    max_tokens: int = 100,
+    temperature: float = 0.8,
+    top_p: float = 0.9,
+    top_k: int = 50,
+    repetition_penalty: float = 1.1,
+) -> str:
+    """Generate text continuation from a prompt."""
+    
+    if not prompt.strip():
+        return "Please enter a prompt to generate text."
+    
+    # Tokenize
+    if USE_HF_TOKENIZER and tokenizer is not None:
+        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+        eos_token_id = tokenizer.eos_token_id
+    else:
+        # Simple fallback - won't work well but better than crashing
+        return "Tokenizer not available. Please ensure transformers is installed."
+    
+    # Generate
+    with torch.no_grad():
+        output_ids = model.generate(
+            input_ids,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            repetition_penalty=repetition_penalty,
+            eos_token_id=eos_token_id,
+        )
+    
+    # Decode
+    if USE_HF_TOKENIZER and tokenizer is not None:
+        generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    else:
+        generated_text = "Decoding not available."
+    
+    return generated_text
+
+
+# Example prompts
+EXAMPLES = [
+    ["The history of artificial intelligence began"],
+    ["In the year 2050, humanity"],
+    ["The most important scientific discovery was"],
+    ["Once upon a time, in a kingdom far away"],
+    ["The universe is vast and"],
+    ["Climate change affects"],
+    ["The theory of relativity states that"],
+    ["In ancient Rome,"],
+]
+
+
+# Create Gradio interface
+with gr.Blocks(title="Tiny-LLM Text Generator") as demo:
+    gr.Markdown("""
+    # 🤖 Tiny-LLM Text Generator
+    
+    A **54 million parameter** language model trained **from scratch** on Wikipedia.
+    
+    This demonstrates that meaningful language models can be trained on consumer hardware!
+    
+    ### Architecture
+    - **Parameters**: 54.93M
+    - **Layers**: 12
+    - **Hidden Size**: 512
+    - **Attention Heads**: 8
+    - **Position Encoding**: RoPE
+    - **Normalization**: RMSNorm
+    - **Activation**: SwiGLU
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            prompt_input = gr.Textbox(
+                label="Prompt",
+                placeholder="Enter your prompt here...",
+                lines=3,
+                value="The history of artificial intelligence began"
+            )
+            
+            with gr.Row():
+                with gr.Column():
+                    max_tokens = gr.Slider(
+                        minimum=10,
+                        maximum=256,
+                        value=100,
+                        step=10,
+                        label="Max New Tokens",
+                    )
+                    temperature = gr.Slider(
+                        minimum=0.1,
+                        maximum=2.0,
+                        value=0.8,
+                        step=0.1,
+                        label="Temperature",
+                        info="Higher = more random"
+                    )
+                
+                with gr.Column():
+                    top_p = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.9,
+                        step=0.05,
+                        label="Top-p (Nucleus Sampling)",
+                    )
+                    top_k = gr.Slider(
+                        minimum=1,
+                        maximum=100,
+                        value=50,
+                        step=5,
+                        label="Top-k",
+                    )
+            
+            repetition_penalty = gr.Slider(
+                minimum=1.0,
+                maximum=2.0,
+                value=1.1,
+                step=0.05,
+                label="Repetition Penalty",
+                info="Higher = less repetition"
+            )
+            
+            generate_btn = gr.Button("✨ Generate", variant="primary", size="lg")
+        
+        with gr.Column(scale=2):
+            output_text = gr.Textbox(
+                label="Generated Text",
+                lines=15,
+                interactive=False,
+            )
+    
+    gr.Markdown("### 📝 Example Prompts")
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=prompt_input,
+    )
+    
+    # Event handlers
+    generate_btn.click(
+        fn=generate_text,
+        inputs=[prompt_input, max_tokens, temperature, top_p, top_k, repetition_penalty],
+        outputs=output_text,
+    )
+    
+    prompt_input.submit(
+        fn=generate_text,
+        inputs=[prompt_input, max_tokens, temperature, top_p, top_k, repetition_penalty],
+        outputs=output_text,
+    )
+    
+    gr.Markdown("""
+    ---
+    ### About This Model
+    
+    **Model**: [jonmabe/tiny-llm-54m](https://huggingface.co/jonmabe/tiny-llm-54m)
+    
+    This is a decoder-only transformer trained from scratch on Wikipedia text.
+    It demonstrates that meaningful language models can be trained on consumer hardware
+    with modest compute budgets (~3 hours on an RTX 5090).
+    
+    #### Training Details
+    - **Training Steps**: 50,000
+    - **Tokens**: ~100M
+    - **Hardware**: NVIDIA RTX 5090 (32GB)
+    - **Training Time**: ~3 hours
+    
+    #### Limitations
+    - Small model size limits knowledge and capabilities
+    - Trained only on Wikipedia - limited domain coverage
+    - May generate factually incorrect information
+    - Not instruction-tuned
+    
+    #### Intended Use
+    - Educational: Understanding transformer training
+    - Experimental: Testing fine-tuning approaches
+    - Research: Lightweight model for NLP experiments
+    """)
+
+
+if __name__ == "__main__":
+    demo.launch()

model.py

@@ -0,0 +1,268 @@
+"""
+TinyLLM Model Architecture
+
+A small transformer language model (~54.93M parameters) trained from scratch.
+Architecture:
+- 12 layers
+- 512 hidden size
+- 8 attention heads
+- 1408 intermediate (FFN)
+- 32000 vocab size
+- 512 max sequence length
+- RoPE position encoding
+- RMSNorm
+- SwiGLU activation
+- Weight tying
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Dict, Any, Optional
+
+
+# Model configuration
+MODEL_CONFIG = {
+    "vocab_size": 32000,
+    "hidden_size": 512,
+    "num_layers": 12,
+    "num_heads": 8,
+    "intermediate_size": 1408,
+    "max_position_embeddings": 512,
+    "dropout": 0.0,
+    "tie_weights": True,
+}
+
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization."""
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary Position Embedding (RoPE)."""
+    def __init__(self, dim: int, max_seq_len: int = 512, base: int = 10000):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self.max_seq_len = max_seq_len
+    
+    def forward(self, seq_len: int, device):
+        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb.cos(), emb.sin()
+
+
+def rotate_half(x):
+    """Rotate half the hidden dims of the input."""
+    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    """Apply rotary positional embeddings to query and key tensors."""
+    cos = cos.unsqueeze(0).unsqueeze(0)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class Attention(nn.Module):
+    """Multi-head attention with RoPE."""
+    def __init__(self, config: Dict[str, Any]):
+        super().__init__()
+        self.hidden_size = config["hidden_size"]
+        self.num_heads = config["num_heads"]
+        self.head_dim = self.hidden_size // self.num_heads
+        
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        
+        self.rotary = RotaryEmbedding(self.head_dim, config["max_position_embeddings"])
+        self.dropout = nn.Dropout(config.get("dropout", 0.0))
+    
+    def forward(self, x, attention_mask=None):
+        B, T, C = x.shape
+        
+        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        
+        cos, sin = self.rotary(T, x.device)
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+        
+        # Scaled dot-product attention
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        
+        # Causal mask
+        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
+        attn_weights.masked_fill_(causal_mask, float('-inf'))
+        
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+        
+        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(x.dtype)
+        attn_weights = self.dropout(attn_weights)
+        
+        out = torch.matmul(attn_weights, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.o_proj(out)
+
+
+class FFN(nn.Module):
+    """Feed-forward network with SwiGLU activation."""
+    def __init__(self, config: Dict[str, Any]):
+        super().__init__()
+        # SwiGLU: w1=gate, w2=down, w3=up
+        self.w1 = nn.Linear(config["hidden_size"], config["intermediate_size"], bias=False)
+        self.w2 = nn.Linear(config["intermediate_size"], config["hidden_size"], bias=False)
+        self.w3 = nn.Linear(config["hidden_size"], config["intermediate_size"], bias=False)
+    
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class TransformerBlock(nn.Module):
+    """Transformer block with pre-norm architecture."""
+    def __init__(self, config: Dict[str, Any]):
+        super().__init__()
+        self.norm1 = RMSNorm(config["hidden_size"])
+        self.attn = Attention(config)
+        self.norm2 = RMSNorm(config["hidden_size"])
+        self.ffn = FFN(config)
+    
+    def forward(self, x, attention_mask=None):
+        x = x + self.attn(self.norm1(x), attention_mask)
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class TinyLLM(nn.Module):
+    """
+    TinyLLM: A small decoder-only transformer language model.
+    
+    Args:
+        config: Dictionary containing model configuration
+    """
+    def __init__(self, config: Dict[str, Any] = None):
+        super().__init__()
+        self.config = config or MODEL_CONFIG
+        
+        self.embed_tokens = nn.Embedding(self.config["vocab_size"], self.config["hidden_size"])
+        self.layers = nn.ModuleList([
+            TransformerBlock(self.config) 
+            for _ in range(self.config["num_layers"])
+        ])
+        self.norm = RMSNorm(self.config["hidden_size"])
+        self.lm_head = nn.Linear(self.config["hidden_size"], self.config["vocab_size"], bias=False)
+        
+        # Tie embeddings if configured
+        if self.config.get("tie_weights", True):
+            self.lm_head.weight = self.embed_tokens.weight
+        
+        # Register causal mask buffer
+        max_len = self.config["max_position_embeddings"]
+        self.register_buffer("causal_mask", 
+            torch.triu(torch.ones(max_len, max_len, dtype=torch.bool), diagonal=1))
+    
+    def forward(self, input_ids, attention_mask=None, labels=None):
+        x = self.embed_tokens(input_ids)
+        
+        for layer in self.layers:
+            x = layer(x, attention_mask)
+        
+        x = self.norm(x)
+        logits = self.lm_head(x)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config["vocab_size"]),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+        
+        return {"logits": logits, "loss": loss}
+    
+    @torch.no_grad()
+    def generate(
+        self, 
+        input_ids: torch.Tensor,
+        max_new_tokens: int = 100,
+        temperature: float = 0.8,
+        top_p: float = 0.9,
+        top_k: int = 50,
+        eos_token_id: Optional[int] = None,
+        repetition_penalty: float = 1.0,
+    ) -> torch.Tensor:
+        """
+        Generate text autoregressively.
+        
+        Args:
+            input_ids: Input token IDs [batch_size, seq_len]
+            max_new_tokens: Maximum number of tokens to generate
+            temperature: Sampling temperature (higher = more random)
+            top_p: Nucleus sampling threshold
+            top_k: Top-k sampling threshold
+            eos_token_id: Token ID that signals end of generation
+            repetition_penalty: Penalty for repeating tokens
+            
+        Returns:
+            Generated token IDs including the prompt
+        """
+        self.eval()
+        
+        for _ in range(max_new_tokens):
+            # Truncate if needed
+            if input_ids.size(1) >= self.config["max_position_embeddings"]:
+                input_ids = input_ids[:, -self.config["max_position_embeddings"]+1:]
+            
+            outputs = self(input_ids)
+            logits = outputs["logits"][:, -1, :]
+            
+            # Apply repetition penalty
+            if repetition_penalty != 1.0:
+                for token_id in set(input_ids[0].tolist()):
+                    logits[0, token_id] /= repetition_penalty
+            
+            # Apply temperature
+            logits = logits / temperature
+            
+            # Top-k filtering
+            if top_k > 0:
+                indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+                logits[indices_to_remove] = float('-inf')
+            
+            # Top-p (nucleus) filtering
+            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')
+            
+            # Sample
+            probs = F.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            input_ids = torch.cat([input_ids, next_token], dim=1)
+            
+            # Check for EOS
+            if eos_token_id is not None and next_token.item() == eos_token_id:
+                break
+        
+        return input_ids

requirements.txt

@@ -0,0 +1,4 @@
+gradio>=4.0.0
+torch>=2.0.0
+transformers>=4.35.0
+huggingface_hub>=0.20.0