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

model_id = "TinyLlama/TinyLlama-1.1B-Chat-v0.3"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

pipe = TextGenerationPipeline(
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=200,
    do_sample=True,
    temperature=0.7,
    top_p=0.95
)

def chat(user_input):
    prompt = f"<|user|>\n{user_input}\n<|assistant|>\n"
    res = pipe(prompt)[0]["generated_text"]
    
    # Cắt bỏ phần prompt để chỉ lấy phần trả lời của assistant
    if "<|assistant|>" in res:
        return res.split("<|assistant|>")[-1].strip()
    else:
        return res.strip()

gr.Interface(
    fn=chat,
    inputs="text",
    outputs="text",
    title="🤖 TinyLlama Chatbot",
    description="Một chatbot nhẹ, chạy mô hình TinyLlama 1.1B"
).launch()