update

app.py

@@ -1,13 +1,11 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from threading import Thread
-import spaces # Quan trọng để dùng ZeroGPU
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import spaces
 
-# 1. Triệu hồi "linh vật" Gemma (Bản 2b-it là hợp lý nhất cho Space)
+# 1. Triệu hồi bộ não Gemma 2B
 model_id = "google/gemma-2b-it"
 
-# Load Tokenizer và Model
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
@@ -15,48 +13,32 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto",
 )
 
-# 2. Tuyệt chiêu xử lý tin nhắn với ZeroGPU
-@spaces.GPU # Cấp quyền dùng GPU tạm thời cho hàm này
-def chat_gemma(message, history):
-    # Chuyển đổi lịch sử chat sang định dạng Gemma hiểu được
+# 2. Hàm xử lý API - Tập trung vào tốc độ và độ chính xác
+@spaces.GPU
+def api_chat(message, history):
+    # Chuyển đổi format cho Gemma
     conversation = []
     for user, assistant in history:
         conversation.extend([{"role": "user", "content": user}, {"role": "model", "content": assistant}])
     conversation.append({"role": "user", "content": message})
 
-    # Tokenize đầu vào
+    # Mã hóa đầu vào
     input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt", add_generation_prompt=True).to(model.device)
     
-    # Thiết lập Streamer để chữ nhảy ra từng chữ cho "ngầu"
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-    
-    generate_kwargs = dict(
-        input_ids=input_ids,
-        streamer=streamer,
+    # Tạo phản hồi (không dùng Streamer để trả về 1 cục JSON cho dễ xử lý ở Web)
+    output_ids = model.generate(
+        input_ids,
         max_new_tokens=1024,
         do_sample=True,
         temperature=0.7,
-        top_k=50,
-        top_p=0.95,
     )
+    
+    # Giải mã và trả kết quả
+    response = tokenizer.decode(output_ids[0][len(input_ids[0]):], skip_special_tokens=True)
+    return response
 
-    # Chạy luồng phụ để generate văn bản
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    partial_message = ""
-    for new_token in streamer:
-        partial_message += new_token
-        yield partial_message
-
-# 3. Khởi tạo giao diện Chatbot (và cũng là cổng API)
-demo = gr.ChatInterface(
-    fn=chat_gemma,
-    title="Gemma Chatbot by Thầy Oáp 🚀",
-    description="Hỏi gì cũng đáp, lú đâu thsông đó!",
-    examples=["Giải thích định luật Newton bằng ngôn ngữ Gen Z", "Viết code Python tạo API"],
-    theme="soft"
-)
+# 3. Khởi tạo interface kiểu "Cổng kết nối"
+demo = gr.ChatInterface(fn=api_chat)
 
 if __name__ == "__main__":
     demo.launch()