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

MODEL_NAME = "LiquidAI/LFM2.5-2.6B"  # or your actual repo id

# -------------------------
# LOAD TOKENIZER + MODEL
# -------------------------

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    dtype=torch.float16  # torch_dtype is deprecated
).to("cuda")  # single-GPU placement, no accelerate needed

model.eval()

# -------------------------
# OPTIONAL GPU TEST
# -------------------------

zero = torch.tensor([0], device="cuda")
print(zero.device)  # should be cuda:0

@spaces.GPU
def greet(n):
    print(zero.device)
    return f"Hello {zero + n} Tensor"

# -------------------------
# CHAT FUNCTION
# -------------------------

def model_chat(message, history):
    # history: list of (user, assistant) tuples
    messages = []
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=True,
            temperature=0.8,
            top_p=0.95,
        )

    # slice off the prompt tokens
    generated = outputs[0][inputs["input_ids"].shape[-1]:]
    response = tokenizer.decode(generated, skip_special_tokens=True)
    return response

# -------------------------
# GRADIO APP
# -------------------------

demo = gr.ChatInterface(
    fn=model_chat,
    title="LiquidAI/LFM2.5-2.6B Chat Demo",
    description="Chat with the LiquidAI/LFM2.5-2.6B model.",
)

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