app.py

import torch
from transformers import AutoTokenizer
from model import TransformerModel  # Replace with your model class
import gradio as gr

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")

# Load the model
def load_model(checkpoint_path):
    # Initialize the model (replace with your model's configuration)
    model = TransformerModel(
        vocab_size=49152,
        hidden_size=576,
        num_hidden_layers=30,
        num_attention_heads=9,
        intermediate_size=1536,
        num_key_value_heads=3,
        max_position_embeddings=2048,
        rms_norm_eps=1e-5,
        hidden_act="silu",
        tie_word_embeddings=True,
        pad_token_id=tokenizer.pad_token_id,
    )
    
    # Load the checkpoint
    checkpoint = torch.load(checkpoint_path, map_location="cpu")
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    return model

# Load the model
model = load_model("checkpoint_5050_quantized.pt")

# Function to generate text
def generate_text(prompt, max_length=50, temperature=1.0, top_k=50):
    # Encode the prompt
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    
    # Generate text
    with torch.no_grad():
        output_ids = model.generate(
            input_ids,
            max_length=max_length,
            temperature=temperature,
            top_k=top_k,
            do_sample=True,
        )
    
    # Decode the generated text
    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return generated_text

# Gradio Interface
def gradio_generate_text(prompt, max_length, temperature, top_k):
    return generate_text(prompt, max_length, temperature, top_k)

# Create the Gradio app
interface = gr.Interface(
    fn=gradio_generate_text,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
        gr.Slider(minimum=10, maximum=200, value=50, label="Max Length"),
        gr.Slider(minimum=0.1, maximum=2.0, value=1.0, label="Temperature"),
        gr.Slider(minimum=1, maximum=100, value=50, label="Top-k Sampling"),
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="Text Generation with SMOL-LM2",
    description="Generate text using the SMOL-LM2 model.",
)

# Launch the app
interface.launch()