import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PretrainedConfig
from huggingface_hub import hf_hub_download
import json
import torch.nn as nn
import torch.nn.functional as F
import math
# Define the model architecture
class SmolLM2Config(PretrainedConfig):
model_type = "smollm2"
def __init__(
self,
vocab_size=49152,
hidden_size=576,
intermediate_size=1536,
num_hidden_layers=30,
num_attention_heads=9,
num_key_value_heads=3,
hidden_act="silu",
max_position_embeddings=2048,
initializer_range=0.041666666666666664,
rms_norm_eps=1e-5,
use_cache=True,
pad_token_id=None,
bos_token_id=0,
eos_token_id=0,
tie_word_embeddings=True,
rope_theta=10000.0,
**kwargs
):
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.hidden_act = hidden_act
self.max_position_embeddings = max_position_embeddings
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
self.rope_theta = rope_theta
super().__init__(
pad_token_id=pad_token_id,
bos_token_id=bos_token_id,
eos_token_id=eos_token_id,
tie_word_embeddings=tie_word_embeddings,
**kwargs
)
# Register the model architecture
from transformers import AutoConfig
AutoConfig.register("smollm2", SmolLM2Config)
class RMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-5):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.eps = eps
def forward(self, x):
variance = x.pow(2).mean(-1, keepdim=True)
x = x * torch.rsqrt(variance + self.eps)
return self.weight * x
def precompute_rope_frequencies(dim: int, max_position_embeddings: int, theta: float = 10000.0):
position = torch.arange(max_position_embeddings).unsqueeze(1) # [seq_len, 1]
div_term = theta ** (torch.arange(0, dim, 2).float() / dim) # [dim/2]
freqs = position / div_term # [seq_len, dim/2]
return freqs
def apply_rotary_embeddings(x: torch.Tensor, freqs: torch.Tensor):
# x shape: [batch, seq_len, heads, head_dim]
# freqs shape: [seq_len, head_dim/2]
x_rot = x.float()
# Reshape freqs to match x's dimensions
freqs = freqs.unsqueeze(0).unsqueeze(2) # [1, seq_len, 1, dim/2]
# Split channels for rotation
x1, x2 = x_rot[..., :x_rot.shape[-1]//2], x_rot[..., x_rot.shape[-1]//2:]
# Apply rotary embeddings
cos = torch.cos(freqs).to(x.device)
sin = torch.sin(freqs).to(x.device)
# Ensure broadcasting dimensions match
cos = cos.expand_as(x1)
sin = sin.expand_as(x1)
# Rotate x1 and x2
x1_rot = x1 * cos - x2 * sin
x2_rot = x2 * cos + x1 * sin
# Concatenate back
return torch.cat([x1_rot, x2_rot], dim=-1).to(x.dtype)
class LlamaAttention(nn.Module):
def __init__(self, config: SmolLM2Config):
super().__init__()
self.hidden_size = config.hidden_size
self.num_heads = config.num_attention_heads
self.num_kv_heads = config.num_key_value_heads
self.head_dim = config.hidden_size // config.num_attention_heads
# Adjust projections to match head dimensions
self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
# Initialize rotary embeddings
self.register_buffer(
"rope_freqs",
precompute_rope_frequencies(
self.head_dim, # Use full head_dim for frequencies
config.max_position_embeddings,
config.rope_theta
),
persistent=False
)
def forward(self, hidden_states, attention_mask=None):
batch_size, seq_length, _ = hidden_states.size()
# Project and reshape
q = self.q_proj(hidden_states).view(batch_size, seq_length, self.num_heads, self.head_dim)
k = self.k_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
v = self.v_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
# Apply rotary embeddings
q = apply_rotary_embeddings(q, self.rope_freqs[:seq_length])
k = apply_rotary_embeddings(k, self.rope_freqs[:seq_length])
# Repeat k/v heads if num_kv_heads < num_heads
if self.num_kv_heads < self.num_heads:
k = k.repeat_interleave(self.num_heads // self.num_kv_heads, dim=2)
v = v.repeat_interleave(self.num_heads // self.num_kv_heads, dim=2)
# Scaled dot-product attention
q = q.transpose(1, 2) # (batch, num_heads, seq_len, head_dim)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
attention_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
if attention_mask is not None:
attention_scores = attention_scores + attention_mask
attention_probs = F.softmax(attention_scores, dim=-1)
context = torch.matmul(attention_probs, v)
context = context.transpose(1, 2).contiguous()
context = context.view(batch_size, seq_length, -1)
return self.o_proj(context)
class LlamaMLP(nn.Module):
def __init__(self, config: SmolLM2Config):
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):
gate = self.act_fn(self.gate_proj(x))
up = self.up_proj(x)
return self.down_proj(gate * up)
class LlamaDecoderLayer(nn.Module):
def __init__(self, config: SmolLM2Config):
super().__init__()
self.self_attn = LlamaAttention(config)
self.mlp = LlamaMLP(config)
self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
def forward(self, hidden_states, attention_mask=None):
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
hidden_states = self.self_attn(hidden_states, attention_mask)
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 SmolLM2ForCausalLM(PreTrainedModel):
config_class = SmolLM2Config
def __init__(self, config):
super().__init__(config)
self.config = config
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
# Add lm_head before weight tying
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
# Initialize weights
self.apply(self._init_weights)
# Tie weights if configured
if config.tie_word_embeddings:
self.lm_head.weight = self.embed_tokens.weight
def _init_weights(self, module):
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
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=self.config.initializer_range)
def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
hidden_states = self.embed_tokens(input_ids)
# Create causal attention mask if none provided
if attention_mask is None:
# Create causal mask
seq_length = input_ids.size(1)
# [batch_size, 1, seq_length, seq_length]
causal_mask = torch.triu(
torch.ones((seq_length, seq_length), dtype=torch.bool, device=input_ids.device),
diagonal=1
).unsqueeze(0).unsqueeze(0)
attention_mask = torch.zeros(
(1, 1, seq_length, seq_length),
dtype=hidden_states.dtype,
device=hidden_states.device
)
attention_mask.masked_fill_(causal_mask, float("-inf"))
for layer in self.layers:
hidden_states = layer(hidden_states, attention_mask)
hidden_states = self.norm(hidden_states)
logits = self.lm_head(hidden_states)
loss = None
if labels is not None:
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))
return logits if loss is None else (loss, logits)
def prepare_inputs_for_generation(self, input_ids, **kwargs):
return {
"input_ids": input_ids,
"attention_mask": kwargs.get("attention_mask", None)
}
def generate(
self,
input_ids,
max_length=100,
temperature=0.7,
top_k=50,
do_sample=True,
num_return_sequences=1,
pad_token_id=None,
eos_token_id=None,
**kwargs
):
cur_len = input_ids.shape[1]
batch_size = input_ids.shape[0]
if max_length < cur_len:
max_length = cur_len
unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
while cur_len < max_length:
# Prepare model inputs
model_inputs = self.prepare_inputs_for_generation(input_ids)
# Forward pass
with torch.no_grad():
outputs = self(**model_inputs)
next_token_logits = outputs[:, -1, :]
# Temperature scaling
if temperature != 1.0 and temperature > 0:
next_token_logits = next_token_logits / temperature
# 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')
# Sample or greedy
if do_sample:
probs = F.softmax(next_token_logits, dim=-1)
next_tokens = torch.multinomial(probs, num_samples=1)
else:
next_tokens = torch.argmax(next_token_logits, dim=-1)
next_tokens = next_tokens.unsqueeze(-1)
# Append next tokens
input_ids = torch.cat([input_ids, next_tokens], dim=-1)
cur_len = input_ids.shape[1]
# Early stopping if all sequences have reached the EOS token
if eos_token_id is not None:
unfinished_sequences = unfinished_sequences.mul(
next_tokens.squeeze(-1).ne(eos_token_id).long()
)
if unfinished_sequences.max() == 0:
break
return input_ids
# Register the model
AutoModelForCausalLM.register(SmolLM2Config, SmolLM2ForCausalLM)
# Cache for model and tokenizer
MODEL = None
TOKENIZER = None
CONFIG = None
def initialize():
global MODEL, TOKENIZER, CONFIG
if MODEL is None:
print("Loading model and tokenizer...")
model_id = "jatingocodeo/SmolLM2"
try:
# Download and load config
print("Loading config...")
config_path = hf_hub_download(repo_id=model_id, filename="config.json")
with open(config_path, 'r') as f:
config_dict = json.load(f)
CONFIG = SmolLM2Config(**config_dict)
# Load tokenizer
print("Loading tokenizer...")
TOKENIZER = AutoTokenizer.from_pretrained(
model_id,
model_max_length=CONFIG.max_position_embeddings,
padding_side="left",
truncation_side="left",
trust_remote_code=True
)
# Make sure we're using the correct special tokens
special_tokens = {
'bos_token': '<|endoftext|>',
'eos_token': '<|endoftext|>',
'unk_token': '<|endoftext|>',
'pad_token': '<|endoftext|>' # Using endoftext as pad token since it's not specified
}
TOKENIZER.add_special_tokens(special_tokens)
# Load model weights
print("Loading model...")
weights_path = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin")
# Initialize model
MODEL = SmolLM2ForCausalLM(CONFIG)
# Resize token embeddings to match tokenizer
MODEL.resize_token_embeddings(len(TOKENIZER))
# Load state dict
state_dict = torch.load(weights_path, map_location="cpu")
MODEL.load_state_dict(state_dict)
# Move model to device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL = MODEL.to(device)
print(f"Model loaded successfully on {device}")
except Exception as e:
print(f"Error initializing: {str(e)}")
raise
def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
# Initialize if not already done
if MODEL is None:
try:
initialize()
except Exception as e:
return f"Failed to initialize model: {str(e)}"
try:
# Process prompt
if not prompt.strip():
return "Please enter a prompt."
# Add BOS token if needed
if not prompt.startswith(TOKENIZER.bos_token):
prompt = TOKENIZER.bos_token + prompt
# Encode prompt
encoded = TOKENIZER.encode_plus(
prompt,
add_special_tokens=True,
return_tensors="pt",
padding=True,
truncation=True,
max_length=CONFIG.max_position_embeddings
)
input_ids = encoded["input_ids"].to(MODEL.device)
attention_mask = encoded["attention_mask"].to(MODEL.device)
# Generate
with torch.no_grad():
outputs = MODEL.generate(
input_ids,
attention_mask=attention_mask,
max_length=min(max_length + len(input_ids[0]), CONFIG.max_position_embeddings),
temperature=max(0.1, min(temperature, 1.0)), # Clamp temperature
top_k=max(1, min(top_k, 100)), # Clamp top_k
do_sample=True if temperature > 0 else False,
num_return_sequences=1,
pad_token_id=TOKENIZER.pad_token_id,
eos_token_id=TOKENIZER.eos_token_id,
)
# Decode and return
generated_text = TOKENIZER.decode(outputs[0], skip_special_tokens=True)
return generated_text.strip()
except Exception as e:
import traceback
traceback.print_exc()
return f"Error during text generation: {str(e)}"
# Create Gradio interface
iface = gr.Interface(
fn=generate_text,
inputs=[
gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=2),
gr.Slider(minimum=10, maximum=200, value=100, step=1, label="Max Length"),
gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top K"),
],
outputs=gr.Textbox(label="Generated Text", lines=5),
title="SmolLM2 Text Generator",
description="Generate text using the fine-tuned SmolLM2 model. Adjust parameters to control the generation.",
examples=[
["Once upon a time", 100, 0.7, 50],
["The quick brown fox", 150, 0.8, 40],
],
allow_flagging="never"
)
# Initialize on startup
try:
initialize()
except Exception as e:
print(f"Warning: Model initialization failed: {str(e)}")
print("Model will be initialized on first request")
if __name__ == "__main__":
iface.launch()