Update app.py

app.py

@@ -6,6 +6,7 @@ from langchain.prompts import PromptTemplate
 from langchain.chains import RetrievalQA, ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
 import warnings
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import os
 from dotenv import load_dotenv
 
@@ -18,7 +19,7 @@ INITIAL_MESSAGE = """Halo! 👋 Saya adalah asisten kesehatan feminacare yang si
 Silakan ajukan pertanyaan apa saja dan saya akan membantu Anda dengan informasi yang akurat."""
 
 # Model configurations
-MODEL_NAME = "SeaLLMs/SeaLLM-7B-v2"
+MODEL_NAME = "SeaLLMs/SeaLLMs-v3-7B-Chat"
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 TOP_K_DOCS = 5
 
@@ -34,15 +35,41 @@ def initialize_models():
   
 def create_llm():
     """Initialize the language model with optimized parameters"""
-    return HuggingFaceHub(
-        repo_id=MODEL_NAME,
-        model_kwargs={
-            "temperature": 0.7,  # Balanced between creativity and accuracy
-            "max_new_tokens": 1024,
-            "top_p": 0.9,
-            "frequency_penalty": 0.5
-        }
+    
+    
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
+    )
+    
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, quantization_config=bnb_config)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    
+    terminators = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|eot_id|>")]
+    
+    text_generation_pipeline = pipeline(
+        model=model,
+        tokenizer=tokenizer,
+        task="text-generation",
+        temperature=0.2,
+        do_sample=True,
+        repetition_penalty=1.1,
+        return_full_text=False,
+        max_new_tokens=200,
+        eos_token_id=terminators,
     )
+    
+    llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
+
+    # return HuggingFaceHub(
+    #     repo_id=MODEL_NAME,
+    #     model_kwargs={
+    #         "temperature": 0.7,  # Balanced between creativity and accuracy
+    #         "max_new_tokens": 1024,
+    #         "top_p": 0.9,
+    #         "frequency_penalty": 0.5
+    #     }
+    # )
+    return llm
 
 # Improved prompt template with better context handling and response structure
 PROMPT_TEMPLATE = """