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

MODEL_NAME = "LiquidAI/LFM2.5-2.6B"

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

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    dtype=torch.float16
).to("cuda")

model.eval()

# -------------------------
# CHAT FUNCTION
# -------------------------
@spaces.GPU
def model_chat(message, history):
    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,
        )

    generated = outputs[0][inputs["input_ids"].shape[-1]:]
    response = tokenizer.decode(generated, skip_special_tokens=True)
    return response

# -------------------------
# FASTAPI ENDPOINT (Spaces auto-serves this)
# -------------------------

app = FastAPI()

@app.post("/hf-chat")
def hf_chat(payload: dict):
    message = payload["message"]
    history = payload.get("history", [])
    return {"response": model_chat(message, history)}

# -------------------------
# GRADIO UI (served on main port)
# -------------------------

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

def main():
    demo.launch()

if __name__ == "__main__":
    main()