Update app.py

app.py

@@ -7,21 +7,27 @@ from langchain.vectorstores import FAISS
 from langchain.chains import RetrievalQA
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import os
-
+import torch
 # Load FLAN-T5 model
 @st.cache_resource
 def load_llm():
-    model_name = "google/flan-t5-base"  # Adjust model size if needed
+    model_name = "tiiuae/falcon-7b"  # Thay bằng tên mô hình bạn chọn
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.bfloat16,  # Giảm bộ nhớ nếu có GPU hỗ trợ
+        trust_remote_code=True,
+        device_map="auto"  # Tự động phân bổ lên GPU/CPU
+    )
     pipe = pipeline(
-        "text2text-generation",
+        "text-generation",
         model=model,
         tokenizer=tokenizer,
-        max_length=512,
-        temperature=0.7,  # Adjust for creativity
+        max_length=512,  # Độ dài tối đa đầu ra
+        temperature=0.7,  # Độ sáng tạo
         top_p=0.95,
-        repetition_penalty=1.15
+        repetition_penalty=1.15,
+        do_sample=True
     )
     return HuggingFacePipeline(pipeline=pipe)