Update app.py

app.py

@@ -5,10 +5,29 @@ 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)
+pipe = TextGenerationPipeline(
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=200,
+    do_sample=True,
+    temperature=0.7,
+    top_p=0.95
+)
 
 def chat(user_input):
-    res = pipe(user_input)[0]["generated_text"]
-    return res[len(user_input):]  # chỉ trả lời, không lặp lại câu hỏi
+    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").launch()
+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()