Adding code for SmolLM2 text generator app

README.md

@@ -9,4 +9,68 @@ app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SmolLM2 Text Generator
+
+This is a Gradio application for generating text using the trained SmolLM2 model. The app allows users to input a text prompt and generate multiple sequences of text based on that prompt. The number of sequences and the length of the generated text can be adjusted using sliders.
+
+## Features
+
+- **Text Generation**: Generate text based on a user-provided prompt using the SmolLM2 model.
+- **Adjustable Length**: Control the length of the generated text.
+- **Multiple Sequences**: Generate multiple sequences of text in one go.
+
+## Requirements
+
+To run this application, you need the following Python packages:
+
+- `torch`
+- `transformers`
+- `gradio`
+
+You can install the required packages using pip:
+
+```bash
+pip install -r requirements.txt
+```
+
+## Usage
+
+1. **Run the App**: Launch the Gradio app by running the following command in your terminal:
+
+   ```bash
+   python app.py
+   ```
+
+2. **Input Prompt**: Enter your desired text prompt in the provided textbox.
+
+3. **Adjust Sliders**:
+   - Use the "Predict Additional Text of Length" slider to set the desired length of the generated text.
+   - Use the "Number of Sequences to Generate" slider to specify how many sequences you want to generate.
+
+4. **Generate Text**: Click the "Generate Text" button to produce the text sequences.
+
+5. **View Output**: The generated sequences will be displayed in the output textbox, each prefixed with "Sequence X:" for clarity.
+
+## Example
+
+- **Prompt**: "Once upon a time"
+- **Number of Sequences**: 2
+
+**Output**:
+```
+Sequence 1:
+Once upon a time, there is a cat ....
+
+Sequence 2:
+Once upon a time in a small village ....
+```
+
+## License
+
+This project is licensed under the MIT License. See the LICENSE file for more details.
+
+## Acknowledgments
+
+- Hugging Face for the Transformers library and model support.
+- Gradio for providing an easy-to-use interface for machine learning applications.
+- The SmolLM2 model for enabling advanced text generation capabilities. 

app.py

@@ -0,0 +1,80 @@
+import gradio as gr
+import torch
+from transformers import AutoTokenizer
+from model import SmolLM2  # Ensure this imports your model correctly
+
+# Load the model and tokenizer
+model_path = "smollm2_final.pt"
+tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")  # Adjust if necessary
+
+# Load model configuration
+model_config = {
+    "bos_token_id": 0,
+    "eos_token_id": 0,
+    "hidden_act": "silu",
+    "hidden_size": 576,
+    "initializer_range": 0.041666666666666664,
+    "intermediate_size": 1536,
+    "is_llama_config": True,
+    "max_position_embeddings": 2048,
+    "num_attention_heads": 9,
+    "num_hidden_layers": 30,
+    "num_key_value_heads": 3,
+    "pad_token_id": None,
+    "pretraining_tp": 1,
+    "rms_norm_eps": 1.0e-05,
+    "rope_interleaved": False,
+    "rope_scaling": None,
+    "rope_theta": 10000.0,
+    "tie_word_embeddings": True,
+    "use_cache": True,
+    "vocab_size": 49152
+}
+
+# Initialize the model with the configuration
+model = SmolLM2(model_config)  # Pass the configuration to the model
+
+# Load the model weights with map_location to handle CPU-only environments
+model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))  # Load the model weights
+model.eval()  # Set the model to evaluation mode
+
+def generate_text(prompt, length, num_sequences):
+    input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"]
+    
+    generated_texts = []
+    for _ in range(num_sequences):
+        generated_sequence = model.generate(
+            input_ids,
+            max_length=length + len(input_ids[0]),  # Adjust for input length
+            pad_token_id=tokenizer.pad_token_id,
+            do_sample=True,
+            temperature=0.8,
+            top_k=50,
+            top_p=0.95
+        )
+        
+        # Decode the generated sequence
+        generated_text = tokenizer.decode(generated_sequence[0], skip_special_tokens=True)
+        generated_texts.append(generated_text)
+    
+    # Format the output
+    formatted_output = "\n\n".join([f"Sequence {i + 1}:\n{text}" for i, text in enumerate(generated_texts)])
+    return formatted_output
+
+# Create Gradio interface
+with gr.Blocks() as app:
+    gr.Markdown("# SmolLM2 Text Generator")
+    prompt_input = gr.Textbox(label="Enter your text prompt", placeholder="Type your prompt here...")
+    length_slider = gr.Slider(minimum=10, maximum=200, label="Predict Additional Text of Length", value=50)
+    num_sequences_slider = gr.Slider(minimum=1, maximum=5, label="Number of Sequences to Generate", value=1, step=1)  # Step set to 1 for integer values
+    generate_button = gr.Button("Generate Text")
+    output_text = gr.Textbox(label="Generated Text", interactive=False)
+
+    generate_button.click(
+        fn=generate_text,
+        inputs=[prompt_input, length_slider, num_sequences_slider],
+        outputs=output_text
+    )
+
+# Launch the app
+app.launch()

model.py

@@ -0,0 +1,245 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def _init_weights(module, std=0.041666666666666664):
+    if isinstance(module, nn.Linear):
+        module.weight.data.normal_(mean=0.0, std=std)
+    elif isinstance(module, nn.Embedding):
+        module.weight.data.normal_(mean=0.0, std=std)
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return x * norm * self.weight
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, theta=10000.0):
+        super().__init__()
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        self.theta = theta
+        
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        
+        t = torch.arange(self.max_position_embeddings).type_as(self.inv_freq)
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos()[None, None, :, :])
+        self.register_buffer("sin_cached", emb.sin()[None, None, :, :])
+
+    def forward(self, x, seq_len=None):
+        if seq_len > self.max_position_embeddings:
+            seq_len = self.max_position_embeddings
+            
+        return (
+            self.cos_cached[:,:,:seq_len,:],
+            self.sin_cached[:,:,:seq_len,:]
+        )
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    # Ensure proper broadcasting
+    cos = cos[:, :, :q.size(2), :]  # [batch, 1, seq_len, dim]
+    sin = sin[:, :, :q.size(2), :]  # [batch, 1, seq_len, dim]
+    
+    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):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config["hidden_size"]
+        self.num_attention_heads = config["num_attention_heads"]
+        self.num_key_value_heads = config["num_key_value_heads"]
+        self.head_dim = self.hidden_size // self.num_attention_heads
+        
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        
+        self.kv_cache = None
+    
+    def forward(self, hidden_states, cos, sin, attention_mask=None, use_cache=False):
+        batch_size, seq_length, _ = hidden_states.shape
+        
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+        
+        # Reshape for attention computation
+        q = q.view(batch_size, seq_length, self.num_attention_heads, self.head_dim)
+        k = k.view(batch_size, seq_length, self.num_key_value_heads, self.head_dim)
+        v = v.view(batch_size, seq_length, self.num_key_value_heads, self.head_dim)
+        
+        # Transpose for attention computation
+        q = q.transpose(1, 2)  # [batch, num_heads, seq_len, head_dim]
+        k = k.transpose(1, 2)  # [batch, num_kv_heads, seq_len, head_dim]
+        v = v.transpose(1, 2)  # [batch, num_kv_heads, seq_len, head_dim]
+        
+        # Apply rotary embeddings
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+        
+        # Repeat k/v heads if num_key_value_heads < num_attention_heads
+        if self.num_key_value_heads != self.num_attention_heads:
+            k = k.repeat_interleave(self.num_attention_heads // self.num_key_value_heads, dim=1)
+            v = v.repeat_interleave(self.num_attention_heads // self.num_key_value_heads, dim=1)
+        
+        # Compute attention
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+            
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        
+        # Compute output
+        output = torch.matmul(attn_weights, v)
+        output = output.transpose(1, 2).contiguous()  # [batch, seq_len, num_heads, head_dim]
+        output = output.view(batch_size, seq_length, -1)
+        
+        return self.o_proj(output)
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.gate_proj = nn.Linear(config["hidden_size"], config["intermediate_size"], bias=False)
+        self.up_proj = nn.Linear(config["hidden_size"], config["intermediate_size"], bias=False)
+        self.down_proj = nn.Linear(config["intermediate_size"], config["hidden_size"], bias=False)
+        self.act_fn = nn.SiLU()
+        
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self_attn = Attention(config)
+        self.mlp = MLP(config)
+        self.input_layernorm = RMSNorm(config["hidden_size"], eps=config["rms_norm_eps"])
+        self.post_attention_layernorm = RMSNorm(config["hidden_size"], eps=config["rms_norm_eps"])
+        
+    def forward(self, hidden_states, cos, sin, attention_mask=None, use_cache=False):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, cos, sin, attention_mask, use_cache)
+        hidden_states = residual + hidden_states
+        
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        
+        return hidden_states
+
+class SmolLM2(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        
+        self.embed_tokens = nn.Embedding(config["vocab_size"], config["hidden_size"])
+        self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config["num_hidden_layers"])])
+        self.norm = RMSNorm(config["hidden_size"], eps=config["rms_norm_eps"])
+        self.rotary_emb = RotaryEmbedding(
+            config["hidden_size"] // config["num_attention_heads"],
+            max_position_embeddings=config["max_position_embeddings"],
+            theta=config.get("rope_theta", 10000.0)
+        )
+        
+        # Initialize weights
+        self.apply(lambda p: _init_weights(p, std=config.get("initializer_range", 0.041666666666666664)))
+        
+    def forward(self, input_ids, attention_mask=None, use_cache=False):
+        hidden_states = self.embed_tokens(input_ids)
+        
+        seq_length = input_ids.shape[1]
+        cos, sin = self.rotary_emb(hidden_states, seq_length)
+        
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, cos, sin, attention_mask, use_cache)
+            
+        hidden_states = self.norm(hidden_states)
+        
+        # Use tied weights for the output projection
+        if self.config.get("tie_word_embeddings", True):
+            logits = F.linear(hidden_states, self.embed_tokens.weight)
+        else:
+            logits = self.lm_head(hidden_states)
+        
+        return logits
+
+    def generate(
+        self, 
+        input_ids, 
+        max_length, 
+        min_length=None,
+        num_return_sequences=1, 
+        pad_token_id=None,
+        do_sample=True,
+        temperature=0.8,
+        top_k=50,
+        top_p=0.95
+    ):
+        self.eval()
+        batch_size = input_ids.shape[0]
+        min_length = min_length if min_length is not None else input_ids.shape[1]
+        
+        # Clear KV cache
+        for layer in self.layers:
+            layer.self_attn.kv_cache = None
+            
+        with torch.no_grad():
+            for _ in range(max_length - input_ids.shape[1]):
+                outputs = self(input_ids, use_cache=True)
+                next_token_logits = outputs[:, -1, :]
+                
+                # Apply temperature
+                next_token_logits = next_token_logits / 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_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+                    next_token_logits[indices_to_remove] = float('-inf')
+                
+                # Sample from the filtered distribution
+                if do_sample:
+                    probs = torch.softmax(next_token_logits, dim=-1)
+                    next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+                else:
+                    next_tokens = torch.argmax(next_token_logits, dim=-1)
+                
+                input_ids = torch.cat([input_ids, next_tokens.unsqueeze(-1)], dim=-1)
+                
+                # Stop if all sequences have hit the pad token
+                if pad_token_id is not None and (next_tokens == pad_token_id).all():
+                    break
+                
+                # Stop if we've reached min_length
+                if input_ids.shape[1] < min_length:
+                    continue
+                    
+        return input_ids 

requirements.txt

@@ -0,0 +1,3 @@
+torch
+transformers
+gradio

smollm2_final.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c96efcee9cd1f94cf2d072647409d1bfce940859d08e89cade4fd48b9502ad2b
+size 269663830