Update app.py

app.py

@@ -1,30 +1,23 @@
 import streamlit as st
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
-from langchain_community.llms import HuggingFaceHub
 from langchain.prompts import PromptTemplate
-from langchain.chains import RetrievalQA, ConversationalRetrievalChain
+from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
 import warnings
-from transformers import pipeline
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import os
 from dotenv import load_dotenv
-from langchain_huggingface import ChatHuggingFace, HuggingFacePipeline
-from langchain_community.llms import HuggingFaceEndpoint
-
+from langchain_huggingface import HuggingFacePipeline
 
 warnings.filterwarnings("ignore")
 load_dotenv()
-api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
 # Constants and configurations
 APP_TITLE = "💊 Asisten Kesehatan Feminacare"
 INITIAL_MESSAGE = """Halo! 👋 Saya adalah asisten kesehatan feminacare yang siap membantu Anda dengan informasi seputar kesehatan wanita. 
 Silakan ajukan pertanyaan apa saja dan saya akan membantu Anda dengan informasi yang akurat."""
 
-# Model configurations
 MODEL_NAME = "SeaLLMs/SeaLLMs-v3-7B-Chat"
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 TOP_K_DOCS = 5
@@ -40,19 +33,21 @@ def initialize_models():
     return vector_store
   
 def create_llm():
-    """Initialize the language model with optimized parameters"""
-    
-    
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
+    """Initialize the language model with auto device mapping"""
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map="auto",
+        trust_remote_code=True
     )
-    
-    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|>")]
+    # Get terminators for the model
+    terminators = [tokenizer.eos_token_id]
+    if hasattr(tokenizer, 'convert_tokens_to_ids'):
+        try:
+            terminators.append(tokenizer.convert_tokens_to_ids("<|eot_id|>"))
+        except:
+            pass
     
     text_generation_pipeline = pipeline(
         model=model,
@@ -66,51 +61,16 @@ def create_llm():
         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
-    #     }
-    # )
-    # llm = HuggingFaceEndpoint(
-    #         repo_id=MODEL_NAME,
-    #         huggingfacehub_api_token = api_token,
-    #         temperature = 0.7,
-    #         max_new_tokens = 1024,
-    #         top_k = 0.9,
-    #     )
-    # llm = HuggingFacePipeline.from_model_id(
-    #     model_id=MODEL_NAME,
-    #     task="text-generation",
-    #     pipeline_kwargs=dict(
-    #         max_new_tokens=512,
-    #         do_sample=False,
-    #         repetition_penalty=1.03,
-    #     ),
-    # )
-    return llm
-    
-    
-    # chat_model = ChatHuggingFace(llm=llm)
+    return HuggingFacePipeline(pipeline=text_generation_pipeline)
 
-# Improved prompt template with better context handling and response structure
 PROMPT_TEMPLATE = """
 Anda adalah asisten kesehatan profesional dengan nama Feminacare. 
 Berikan informasi yang akurat, jelas, dan bermanfaat berdasarkan konteks yang tersedia.
-
 Context yang tersedia:
 {context}
-
 Chat historyt:
 {chat_history}
-
 Question: {question}
-
 Instruksi untuk menjawab:
 1. Berikan jawaban yang LENGKAP dan TERSTRUKTUR
 2. Selalu sertakan SUMBER informasi dari konteks yang diberikan
@@ -118,7 +78,6 @@ Instruksi untuk menjawab:
 4. Gunakan bahasa yang mudah dipahami
 5. Jika relevan, berikan poin-poin penting menggunakan format yang rapi
 6. Akhiri dengan anjuran untuk konsultasi dengan tenaga kesehatan jika diperlukan
-
 Answer:
 """
 
@@ -137,14 +96,10 @@ def setup_qa_chain(vector_store):
     
     return ConversationalRetrievalChain.from_llm(
         llm=create_llm(),
-        retriever=vector_store.as_retriever(
-            # search_type="mmr",  # Maximum Marginal Relevance for better diversity
-            # search_kwargs={"k": TOP_K_DOCS}
-        ),
+        retriever=vector_store.as_retriever(),
         memory=memory,
-        # combine_docs_chain_kwargs={"prompt": custom_prompt},
+        combine_docs_chain_kwargs={"prompt": custom_prompt},
         return_source_documents=True,
-        # return_generated_question=True,
     )
 
 def initialize_session_state():
@@ -199,8 +154,8 @@ def create_ui():
 def handle_user_input(prompt):
     """Handle user input and generate response"""
     with st.spinner("Sedang menyiapkan jawaban..."):
-            response = st.session_state.qa_chain({"question": prompt})
-            return response["answer"]
+        response = st.session_state.qa_chain({"question": prompt})
+        return response["answer"]
 
 def main():
     initialize_session_state()