app.py

import gradio as gr
import torch
from beeper_model import BeeperRoseGPT, generate
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file as load_safetensors

# ----------------------------
# 🔧 Load Model and Tokenizer
# ----------------------------
config = {
    "context": 512,
    "vocab_size": 8192,
    "dim": 512,
    "n_heads": 8,
    "n_layers": 6,
    "mlp_ratio": 4.0,
    "temperature": 0.9,
    "top_k": 40,
    "top_p": 0.9,
    "repetition_penalty": 1.1,
    "presence_penalty": 0.6,
    "frequency_penalty": 0.0,
    "resid_dropout": 0.1,  # Add these for model init
    "dropout": 0.0,
    "grad_checkpoint": False,
    "tokenizer_path": "beeper.tokenizer.json"
}

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load weights from Hugging Face repo
repo_id = "AbstractPhil/beeper-rose-tinystories-6l-512d-ctx512"
model_file = hf_hub_download(repo_id=repo_id, filename="beeper_rose_final.safetensors")
tokenizer_file = hf_hub_download(repo_id=repo_id, filename="tokenizer.json")

# Initialize model
infer = BeeperRoseGPT(config).to(device)

# Load safetensors properly
state_dict = load_safetensors(model_file, device=str(device))
infer.load_state_dict(state_dict)
infer.eval()

# Load tokenizer
tok = Tokenizer.from_file(tokenizer_file)

# ----------------------------
# 💬 Gradio Chat Wrapper
# ----------------------------
def beeper_reply(message, history, temperature, top_k, top_p):
    # Build conversation context
    prompt_parts = []
    for h in history:
        if h[0]:  # User message exists
            prompt_parts.append(f"User: {h[0]}")
        if h[1]:  # Assistant response exists
            prompt_parts.append(f"Beeper: {h[1]}")
    
    # Add current message
    prompt_parts.append(f"User: {message}")
    prompt_parts.append("Beeper:")
    
    prompt = "\n".join(prompt_parts)
    
    # Generate response
    response = generate(
        model=infer,
        tok=tok,
        cfg=config,
        prompt=prompt,
        max_new_tokens=128,
        temperature=temperature,
        top_k=int(top_k),
        top_p=top_p,
        repetition_penalty=config["repetition_penalty"],
        presence_penalty=config["presence_penalty"],
        frequency_penalty=config["frequency_penalty"],
        device=device,
        detokenize=True
    )
    
    # Clean up response - remove the prompt part if it's included
    if response.startswith(prompt):
        response = response[len(prompt):].strip()
    
    return response

# ----------------------------
# 🖼️ Interface
# ----------------------------
demo = gr.ChatInterface(
    beeper_reply,
    additional_inputs=[
        gr.Slider(0.1, 1.5, value=0.9, step=0.1, label="Temperature"),
        gr.Slider(1, 100, value=40, step=1, label="Top-k"),
        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p"),
    ],
    chatbot=gr.Chatbot(label="Chat with Beeper 🤖"),
    title="Beeper - A Rose-based Tiny Language Model",
    description="Hello! I'm Beeper, a small language model trained with love and care. Please be patient with me - I'm still learning! 💕",
    examples=[
        ["Hello Beeper! How are you today?"],
        ["Can you tell me a story about a robot?"],
        ["What do you like to do for fun?"],
    ],
    theme=gr.themes.Soft(),
)

if __name__ == "__main__":
    demo.launch()