import spaces
import torch
import gradio as gr
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "IFM/K2-Horizon-0.9B-GGUF"


@spaces.GPU(duration=60)
def load_model():
    """Load the model and tokenizer with proper config handling."""
    try:
        config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True)
    except ValueError:
        # Model config doesn't have model_type, set it manually
        config = AutoConfig.for_model("k2_horizon")
        config.is_encoder_decoder = False
    
    if not hasattr(config, "model_type") or config.model_type is None:
        config.model_type = "k2_horizon"
    
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID, 
        trust_remote_code=True, 
        use_fast=False
    )
    
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16,
        device_map="auto",
        config=config,
        trust_remote_code=True,
    )
    
    return model, tokenizer


model, tokenizer = load_model()


@spaces.GPU(duration=60)
def generate(message, history):
    """Generate a response from the K2-Horizon model."""
    # Format messages for the model
    if history is None:
        history = []
    
    history.append((message, ""))
    
    # Build chat template - K2-Horizon uses a specific format
    messages = []
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})
    
    messages.append({"role": "user", "content": message})
    
    # Tokenize using the tokenizer
    input_ids = tokenizer.apply_chat_template(
        messages, 
        tokenize=True, 
        add_generation_prompt=True, 
        return_tensors="pt"
    ).to(model.device)
    
    # Generate
    with torch.no_grad():
        outputs = model.generate(
            input_ids,
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
        )
    
    # Decode response
    response = tokenizer.decode(
        outputs[0][input_ids.shape[1]:], skip_special_tokens=True
    )
    
    history[-1] = (message, response)
    return response


demo = gr.ChatInterface(
    fn=generate,
    title="K2-Horizon-0.9B-GGUF",
    description="Chat with the K2-Horizon 0.9B reasoning model (GGUF quantized).",
    examples=[
        ["Who are you?"],
        ["What is 84 * 3 / 2?"],
        ["Tell me an interesting fact about the universe!"],
        ["Explain quantum computing in simple terms."],
    ],
    cache_examples=True,
)

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