import gradio as gr
import torch
import json
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download
from ModelArchitecture import Transformer, ModelConfig, generate
from safetensors.torch import load_file

# Load model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
REPO_ID = "VirtualInsight/Lumen"

model_path = hf_hub_download(repo_id=REPO_ID, filename="model.safetensors")
tokenizer_path = hf_hub_download(repo_id=REPO_ID, filename="tokenizer.json")
config_path = hf_hub_download(repo_id=REPO_ID, filename="config.json")

tokenizer = Tokenizer.from_file(tokenizer_path)
with open(config_path) as f:
    config = ModelConfig(**json.load(f))

model = Transformer(config).to(device)
model.load_state_dict(load_file(model_path, device=str(device)), strict=False)
model.eval()

@torch.no_grad()
def generate_text(prompt, max_tokens=100, temperature=0.7, top_p=0.9):
    input_ids = torch.tensor(tokenizer.encode(prompt).ids).unsqueeze(0).to(device)
    output_ids = generate(model, input_ids, max_tokens, temperature, top_p=top_p, device=device)
    return tokenizer.decode(output_ids[0, input_ids.size(1):].cpu().tolist())

# Gradio Interface
demo = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Once upon a time...", lines=3),
        gr.Slider(10, 500, value=100, label="Max Tokens"),
        gr.Slider(0.1, 2.0, value=0.7, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.9, label="Top-p"),
    ],
    outputs=gr.Textbox(label="Generated Text", lines=10),
    title="LumenBase Language Model",
    description="Generate text using the Lumen language model",
)

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