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

# =========================
# MODEL (CPU)
# =========================
torch.set_default_device("cpu")

model = AutoModelForCausalLM.from_pretrained(
    "microsoft/phi-1",
    torch_dtype=torch.float32,   # important pour CPU
    trust_remote_code=True
)

tokenizer = AutoTokenizer.from_pretrained(
    "microsoft/phi-2",
    trust_remote_code=True
)

# =========================
# CHAT
# =========================
def chat(message):
    inputs = tokenizer(
        message,
        return_tensors="pt",
        return_attention_mask=False
    )

    outputs = model.generate(
        **inputs,
       
        do_sample=True,
        temperature=0.7
    )

    return tokenizer.batch_decode(outputs)[0]


# =========================
# UI + API
# =========================
with gr.Blocks() as demo:

    gr.Markdown("# 🤖 Phi-2 CPU Chat")

    inp = gr.Textbox(label="Message")
    out = gr.Textbox(label="Réponse")

    btn = gr.Button("Send")

    btn.click(chat, inp, out)

    # ✅ endpoint API auto
    gr.Interface(
        fn=chat,
        inputs=gr.Textbox(),
        outputs=gr.Textbox()
    )

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