Update app.py

app.py

@@ -4,6 +4,26 @@ from typing import List
 from huggingface_hub import login
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+from transformers import AutoTokenizer as SummarizerTokenizer, AutoModelForSeq2SeqLM
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Summarization model
+summarizer_model_id = "facebook/bart-large-cnn"
+summarizer_tokenizer = SummarizerTokenizer.from_pretrained(summarizer_model_id)
+summarizer_model = AutoModelForSeq2SeqLM.from_pretrained(summarizer_model_id)
+summarizer_model.to(device)
+
+def summarize_text(text: str, max_length=150) -> str:
+    inputs = summarizer_tokenizer([text], return_tensors="pt", max_length=1024, truncation=True).to(device)
+    summary_ids = summarizer_model.generate(
+        inputs['input_ids'],
+        num_beams=4,
+        max_length=max_length,
+        early_stopping=True
+    )
+    summary = summarizer_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return summary
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 if HF_TOKEN:
@@ -17,13 +37,17 @@ model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True,
     token=HF_TOKEN
 )
-device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
+# --- GIỮ LẠI CHỈ 1 HÀM build_prompt, ĐÃ BỔ SUNG SUMMARIZATION ---
 def build_prompt(prompt: str, histories: List[str], new_message: str) -> str:
     prompt_text = prompt.strip() + "\n" if prompt else ""
-    if histories:
-        prompt_text += "\n".join(histories) + "\n"
+    histories_text = "\n".join(histories) if histories else ""
+    # Tóm tắt nếu quá dài (tùy chỉnh ngưỡng này)
+    if len(histories_text) > 3000:
+        histories_text = summarize_text(histories_text, max_length=180)
+    if histories_text:
+        prompt_text += histories_text + "\n"
     prompt_text += f"User: {new_message}\nAI:"
     return prompt_text