import torch import gradio as gr from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PretrainedConfig import torch.nn as nn import torch.nn.functional as F import math 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 ) 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 class LlamaAttention(nn.Module): def __init__(self, config): 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 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) 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) # Repeat k/v heads if needed 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) # Transpose for attention q = q.transpose(1, 2) # (batch, num_heads, seq_len, head_dim) k = k.transpose(1, 2) # (batch, num_heads, seq_len, head_dim) v = v.transpose(1, 2) # (batch, num_heads, seq_len, head_dim) # Calculate attention scores scale = 1.0 / math.sqrt(self.head_dim) scores = torch.matmul(q, k.transpose(-2, -1)) * scale # (batch, num_heads, seq_len, seq_len) # Apply attention mask if provided if attention_mask is not None: # Ensure mask is broadcastable if attention_mask.dim() == 2: attention_mask = attention_mask.unsqueeze(1).unsqueeze(1) # (batch, 1, 1, seq_len) scores = scores + attention_mask # Apply softmax and dropout attention_weights = F.softmax(scores, dim=-1) # Apply attention to values output = torch.matmul(attention_weights, v) # (batch, num_heads, seq_len, head_dim) # Reshape and project back output = output.transpose(1, 2).contiguous() # (batch, seq_len, num_heads, head_dim) output = output.view(batch_size, seq_length, -1) # (batch, seq_len, hidden_size) output = self.o_proj(output) return output class LlamaMLP(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): 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): 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 _no_split_modules = ["LlamaDecoderLayer"] 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) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) if config.tie_word_embeddings: self.lm_head.weight = self.embed_tokens.weight def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs): return_dict = return_dict if return_dict is not None else self.config.use_return_dict hidden_states = self.embed_tokens(input_ids) # Create causal attention mask if none provided if attention_mask is None: attention_mask = torch.triu( torch.ones((input_ids.size(1), input_ids.size(1)), dtype=torch.bool, device=input_ids.device), diagonal=1 ) attention_mask = attention_mask.unsqueeze(0).unsqueeze(0) attention_mask = attention_mask * -1e4 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)) if return_dict: from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions return CausalLMOutputWithCrossAttentions( loss=loss, logits=logits, past_key_values=None, hidden_states=None, attentions=None, cross_attentions=None, ) return (loss, logits) if loss is not None else logits def prepare_inputs_for_generation(self, input_ids, **kwargs): return { "input_ids": input_ids, "attention_mask": kwargs.get("attention_mask", None) } # Register the model architecture from transformers import AutoConfig AutoConfig.register("smollm2", SmolLM2Config) AutoModelForCausalLM.register(SmolLM2Config, SmolLM2ForCausalLM) # Cache for model and tokenizer MODEL = None TOKENIZER = None def initialize(): global MODEL, TOKENIZER if MODEL is None: print("Loading model and tokenizer...") model_id = "jatingocodeo/SmolLM2" try: # Load tokenizer print("\n1. Loading tokenizer...") TOKENIZER = AutoTokenizer.from_pretrained(model_id) print("✓ Tokenizer loaded successfully") # Add special tokens if needed special_tokens = { 'pad_token': '[PAD]', 'eos_token': '', 'bos_token': '' } num_added = TOKENIZER.add_special_tokens(special_tokens) print(f"✓ Added {num_added} special tokens") # Load model print("\n2. Loading model...") MODEL = AutoModelForCausalLM.from_pretrained( model_id, trust_remote_code=True, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, low_cpu_mem_usage=True ) # Move model to appropriate device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") MODEL = MODEL.to(device) print(f"✓ Model loaded successfully and moved to {device}") except Exception as e: print(f"Error initializing model: {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: initialize() try: # Process prompt if not prompt.strip(): return "Please enter a prompt." if not prompt.startswith(TOKENIZER.bos_token): prompt = TOKENIZER.bos_token + prompt # Encode prompt input_ids = TOKENIZER.encode(prompt, return_tensors="pt", truncation=True, max_length=2048) input_ids = input_ids.to(MODEL.device) # Generate with torch.no_grad(): output_ids = MODEL.generate( input_ids, max_length=min(max_length + len(input_ids[0]), 2048), temperature=temperature, top_k=top_k, do_sample=True, pad_token_id=TOKENIZER.pad_token_id, eos_token_id=TOKENIZER.eos_token_id, num_return_sequences=1 ) # Decode and return generated_text = TOKENIZER.decode(output_ids[0], skip_special_tokens=True) return generated_text.strip() except Exception as e: return f"Error generating text: {str(e)}" # Initialize on startup initialize() # 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. - Max Length: Controls the length of generated text - Temperature: Controls randomness (higher = more creative) - Top K: Controls diversity of word choices""", examples=[ ["Once upon a time", 100, 0.7, 50], ["The quick brown fox", 150, 0.8, 40], ["In a galaxy far far away", 200, 0.9, 30], ], allow_flagging="never" ) if __name__ == "__main__": iface.launch()