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")

def load_quantized_model(checkpoint_path):
    # Define the model architecture
    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,
    )
    
    # Apply dynamic quantization to the embedding layer
    model.embed_tokens = torch.quantization.quantize_dynamic(
        model.embed_tokens, {torch.nn.Embedding}, dtype=torch.qint8
    )
    
    # Apply static quantization to the rest of the model
    model.qconfig = torch.quantization.default_qconfig
    model = torch.quantization.prepare(model, inplace=False)
    model = torch.quantization.convert(model, inplace=False)
    
    # Load the quantized checkpoint
    checkpoint = torch.load(checkpoint_path, map_location="cpu")
    model.load_state_dict(checkpoint["model_state_dict"])
    
    model.eval()
    return model


import gradio as gr

# Load the quantized model
model = load_quantized_model("checkpoint_quantized.pt")

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

# Gradio Interface
interface = gr.Interface(
    fn=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 Quantized SMOL-LM2",
    description="Generate text using a quantized version of the SMOL-LM2 model.",
)

# Launch the app
interface.launch()