Update generator.py

generator.py

@@ -1,21 +1,27 @@
-from model_loader import load_model
+import torch
+from transformers import TextIteratorStreamer
+from threading import Thread
+from model_loader import model, tokenizer
 
-# load 1 lần để tái sử dụng
-pipe = load_model()
+# Prompt ngầm
+SYSTEM_PROMPT = "Bạn là chatbot chuyên tóm tắt và tổng hợp ý chính từ nội dung người dùng đưa ra."
 
-SYSTEM_PROMPT = (
-    "Bạn là một chatbot chuyên TÓM TẮT và TỔNG HỢP ý chính. "
-    "Nhiệm vụ: nhận câu hỏi hoặc đoạn văn bản và trả lời ngắn gọn, "
-    "liệt kê các điểm quan trọng nhất.\n\n"
-)
+def generate_text(prompt, max_new_tokens=1024):
+    inputs = tokenizer(SYSTEM_PROMPT + "\nUser: " + prompt + "\nBot:", return_tensors="pt").to("cpu")
 
-def generate_text(user_input: str, max_new_tokens: int = 128):
-    prompt = SYSTEM_PROMPT + "Người dùng hỏi: " + user_input + "\n\nTóm tắt & tổng hợp:"
-    outputs = pipe(
-        prompt,
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = dict(
+        **inputs,
         max_new_tokens=max_new_tokens,
-        do_sample=False,      # greedy decoding → nhanh và ổn định hơn
-        temperature=0.6,
-        top_p=0.9
+        do_sample=False,   # Greedy decoding để nhanh hơn
+        temperature=0.7,
+        top_p=0.9,
+        streamer=streamer
     )
-    return outputs[0]["generated_text"]
+
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # Trả về stream kết quả
+    for new_text in streamer:
+        yield new_text