import spaces
import torch
import gradio as gr
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

MODEL_REPO = "ClokAI/ci-base"

print("Loading model...")

from clokai import CiModel, CiConfig
from transformers import AutoTokenizer

config = CiConfig()
model = CiModel(config)

weights_path = hf_hub_download(repo_id=MODEL_REPO, filename="model.safetensors")
sd = load_file(weights_path)
model.load_state_dict(sd, strict=False)
model.eval()

tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
print("Model ready!")

@spaces.GPU
def chat(message, history, temperature, max_tokens):
    device = next(model.parameters()).device
    formatted = f"<s> User: {message} Assistant:"
    inputs = tokenizer(formatted, return_tensors="pt").to(device)
    
    generated = []
    input_ids = inputs["input_ids"]
    
    for _ in range(max_tokens):
        with torch.no_grad():
            logits = model(input_ids).logits
        probs = torch.softmax(logits[:, -1, :] / temperature, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        if next_token.item() == tokenizer.eos_token_id:
            break
        generated.append(next_token.item())
        input_ids = torch.cat([input_ids, next_token], dim=-1)
    
    return tokenizer.decode(generated, skip_special_tokens=True)

demo = gr.ChatInterface(
    fn=chat,
    title="ClokAI ci-base Test",
    description="Test ci-base model - safetensors + clokai library",
    additional_inputs=[
        gr.Slider(0.1, 2.0, value=0.6, step=0.1, label="Temperature"),
        gr.Slider(50, 500, value=200, step=50, label="Max Tokens"),
    ],
)

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