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

model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_id)

generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    do_sample=True,
    top_p=0.9,
    temperature=0.7,
)

def respond(message, history):
    prompt = ""
    for user, bot in history:
        prompt += f"<|user|>: {user}\n<|assistant|>: {bot}\n"
    prompt += f"<|user|>: {message}\n<|assistant|>:"

    output = generator(prompt, max_new_tokens=256)[0]["generated_text"]
    answer = output.split("<|assistant|>:")[-1].strip()
    return answer

gr.ChatInterface(respond).launch(share=True)  # ou server_name="0.0.0.0" pour local