ChatBot / app.py
lea97338's picture
Update app.py
e618b15 verified
Raw
History Blame Contribute Delete
1.31 kB
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()