Update app.py

app.py

@@ -6,39 +6,44 @@ from langchain.embeddings import SentenceTransformerEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 
-# --- Load mô hình ngôn ngữ ---
+# --- 1. Load Mô Hình TinyLlama hoặc Mistral ---
 @st.cache_resource
 def load_llm():
-    model_name = "google/flan-ul2"  # Có thể thay bằng google/flan-ul2 hoặc mistralai/Mistral-7B-Instruct-v0.2 nếu có GPU
-    
+    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # Thay bằng "mistralai/Mistral-7B-Instruct-v0.2" nếu có GPU
+
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(
+    
+    if tokenizer.pad_token is None:
+        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+
+    model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float32,
+        torch_dtype=torch.float32,  # Trên CPU nên dùng float32
         device_map="auto"
     )
-    
+
     pipe = pipeline(
-    "text2text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=512,
-    temperature=0.75,         # Tăng tính sáng tạo (tự nhiên hơn)
-    top_p=0.92,               # Kết hợp với temperature
-    top_k=40,
-    repetition_penalty=1.25,  # Tránh lặp từ ngữ
-    num_beams=4,              # Giúp câu mượt hơn nếu không sampling
-    early_stopping=True,
-    do_sample=True            # Bật chế độ sampling để tránh cứng nhắc
-)
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.9,
+        top_k=50,
+        repetition_penalty=1.2,
+        do_sample=True,
+        eos_token_id=tokenizer.eos_token_id,
+        truncation=True,
+        return_full_text=False
+    )
 
     return HuggingFacePipeline(pipeline=pipe)
 
 
-# --- Xử lý file PDF ---
+# --- 2. Xử lý file PDF ---
 def process_pdf(pdf_path):
     loader = PyPDFLoader(pdf_path)
     documents = loader.load()
@@ -55,8 +60,8 @@ def process_pdf(pdf_path):
     return vectorstore
 
 
-# --- Tiền xử lý prompt và hậu xử lý câu trả lời ---
-template = """あなたは親しみやすく丁寧なアシスタントです。以下の文書情報をもとに、質問に自然で分かりやすい日本語で回答してください。
+# --- 3. Prompt Template tiếng Nhật (tự nhiên) ---
+template = """<s>[INST]あなたは親しみやすく丁寧なアシスタントです。以下の文書情報をもとに、質問に自然で分かりやすい日本語で回答してください。
 
 - 回答はできるだけ口語的で柔らかい表現を使ってください。
 - 理由や例を交えて説明すると良いでしょう。
@@ -66,37 +71,31 @@ template = """あなたは親しみやすく丁寧なアシスタントです。
 {context}
 
 質問: {question}
-回答:"""
+回答: [/INST]"""
 
-QA_PROMPT = PromptTemplate(
-    template=template,
-    input_variables=["context", "question"]
-)
+QA_PROMPT = PromptTemplate(template=template, input_variables=["context", "question"])
 
 
+# --- 4. Hàm hậu xử lý câu trả lời ---
 def postprocess_answer(answer):
     answer = answer.strip()
-
-    # Loại bỏ các cụm không mong muốn
+    
     for phrase in ["Answer:", "答え:", "回答:", "The answer is", "Based on the context"]:
         answer = answer.replace(phrase, "").strip()
 
-    # Viết hoa chữ cái đầu tiên (nếu cần thiết trong tiếng Nhật)
-    # if answer and len(answer) > 0:
-    #     answer = answer[0].upper() + answer[1:]
+    if answer and len(answer) > 0:
+        answer = answer[0].upper() + answer[1:]
 
-    # # Thêm dấu chấm cuối câu nếu thiếu
-    # if answer and answer[-1] not in "。.?！":
-    #     answer += "。"
+    if answer and answer[-1] not in "。.?！":
+        answer += "。"
 
-    # # Nếu câu quá ngắn hoặc vô nghĩa, phản hồi mềm mại
-    # if len(answer.split()) < 4:
-    #     answer = "資料にはその件についての詳細な記載が見受けられませんが、以下のように推測されます：" + answer
+    if len(answer.split()) < 4:
+        answer = "資料にはその件についての詳細な記載が見受けられませんが、以下のように推測されます：" + answer
 
     return answer
 
 
-# --- Giao diện chính của ứng dụng ---
+# --- 5. Giao diện chính của ứng dụng ---
 def main():
     st.set_page_config(page_title="PDFアシスタント", page_icon="📘")
     st.title("PDFアシスタント 🤖")
@@ -112,9 +111,6 @@ def main():
             vectorstore = process_pdf("temp.pdf")
         
         llm = load_llm()
-        response = llm("東京の人口はどのくらいですか？")
-        st.success(response)
-        print("LLM response:", response)
         
         qa_chain = RetrievalQA.from_chain_type(
             llm=llm,
@@ -133,9 +129,8 @@ def main():
             with st.spinner("回答を生成中..."):
                 try:
                     result = qa_chain({"question": query})
-                    answer = result["result"]
-                    # answer = postprocess_answer(raw_answer)
-                    
+                    raw_answer = result["result"]
+                    answer = postprocess_answer(raw_answer)
                     
                     st.markdown("### 回答")
                     st.success(answer)