import gradio as gr
import torch
from beeper_model import BeeperRoseGPT, generate  # assumed modular split
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download

# ----------------------------
# 🔧 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,
    "tokenizer_path": "beeper.tokenizer.json"
}

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

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

infer = BeeperRoseGPT(config).to(device)
infer.load_state_dict(torch.load(model_file, map_location=device))
infer.eval()
tok = Tokenizer.from_file(tokenizer_file)

# ----------------------------
# 💬 Gradio Chat Wrapper
# ----------------------------
def beeper_reply(message, history, temperature, top_k, top_p):
    prompt = "\n".join([f"User: {h[0]}\nBeeper: {h[1]}" for h in history if h[0] and h[1]])
    prompt += f"\nUser: {message}\nBeeper:"

    out = generate(
        model=infer,
        tok=tok,
        cfg=config,
        prompt=prompt,
        max_new_tokens=128,
        temperature=temperature,
        top_k=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
    )
    yield out

# ----------------------------
# 🖼️ 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="Hello I'm Beeper (Rose-based LLM)! Please be friendly I don't know very much yet!")
)

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