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

# Load the model and tokenizer
model_name = "Smilyai-labs/Sam-reason-S2"  # or your local path
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.eval()

# Chat function
def chat_with_model(user_input, history):
    if history is None:
        history = []

    # Build conversation string for context
    conversation = ""
    for msg in history:
        conversation += f"User: {msg['content']}\nSam:"
        if msg['role'] == "assistant":
            conversation += f" {msg['content']}\n"

    conversation += f"User: {user_input}\nSam:"

    # Encode and generate
    inputs = tokenizer(conversation, return_tensors="pt", truncation=True, max_length=1024)
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_new_tokens=150,
            do_sample=True,
            top_k=50,
            top_p=0.95,
            temperature=0.7,
            pad_token_id=tokenizer.eos_token_id
        )

    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response_text = decoded.split("Sam:")[-1].strip()

    # Add new message to chat history
    history.append({"role": "user", "content": user_input})
    history.append({"role": "assistant", "content": response_text})

    return "", history

# Gradio UI
def create_chatbot_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# 💬 Chat with **Sam** (SmilyAI's Reasoning LLM 2nd generation)")
        chatbot = gr.Chatbot(label="Chat", type="messages")
        user_input = gr.Textbox(placeholder="Type your message...", show_label=False)
        send_btn = gr.Button("Send")

        send_btn.click(
            chat_with_model,
            inputs=[user_input, chatbot],
            outputs=[user_input, chatbot]
        )

    return demo

# Launch
demo = create_chatbot_interface()
demo.launch()