Update app.py

app.py

@@ -1,26 +1,27 @@
 import gradio as gr
 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 HuggingFacePipeline
 
 warnings.filterwarnings("ignore")
 load_dotenv()
 
 # 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."""
+TITLE = "💊 Asisten Kesehatan Feminacare"
+DESCRIPTION = """
+# 💊 Asisten Kesehatan Feminacare
+Asisten digital ini dirancang untuk membantu Anda berkonsultasi tentang kesehatan wanita.
+
+*Catatan: Informasi yang diberikan bersifat umum. Selalu konsultasikan dengan tenaga kesehatan untuk saran yang lebih spesifik.*
+"""
 
-# Model configurations
 MODEL_NAME = "SeaLLMs/SeaLLMs-v3-7B-Chat"
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 TOP_K_DOCS = 5
@@ -36,23 +37,20 @@ 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",
-        torch_dtype=torch.float16,
-        quantization_config=bnb_config
+        trust_remote_code=True
     )
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     
-    terminators = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|eot_id|>")]
+    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,
@@ -64,7 +62,6 @@ def create_llm():
         return_full_text=False,
         max_new_tokens=200,
         eos_token_id=terminators,
-        device_map="auto"
     )
     
     return HuggingFacePipeline(pipeline=text_generation_pipeline)
@@ -74,7 +71,7 @@ Anda adalah asisten kesehatan profesional dengan nama Feminacare.
 Berikan informasi yang akurat, jelas, dan bermanfaat berdasarkan konteks yang tersedia.
 Context yang tersedia:
 {context}
-Chat history:
+Chat historyt:
 {chat_history}
 Question: {question}
 Instruksi untuk menjawab:
@@ -89,91 +86,98 @@ Answer:
 
 class HealthAssistant:
     def __init__(self):
-        self.vector_store = initialize_models()
+        vector_store = initialize_models()
         self.memory = ConversationBufferMemory(
             memory_key="chat_history",
             return_messages=True,
             output_key='answer'
         )
-        self.qa_chain = self.setup_qa_chain()
-
-    def setup_qa_chain(self):
-        """Set up the QA chain with improved configuration"""
+        
         custom_prompt = PromptTemplate(
             template=PROMPT_TEMPLATE,
             input_variables=["context", "question", "chat_history"]
         )
         
-        return ConversationalRetrievalChain.from_llm(
+        self.qa_chain = ConversationalRetrievalChain.from_llm(
             llm=create_llm(),
-            retriever=self.vector_store.as_retriever(),
+            retriever=vector_store.as_retriever(),
             memory=self.memory,
+            combine_docs_chain_kwargs={"prompt": custom_prompt},
             return_source_documents=True,
         )
 
     def respond(self, message, history):
-        """Generate response for user input"""
-        if not message:
-            return ""
-        
+        """Process the message and return a response"""
         response = self.qa_chain({"question": message})
         return response["answer"]
-
+    
     def clear_history(self):
-        """Clear conversation history"""
+        """Clear the conversation memory"""
         self.memory.clear()
-        return []
+        return None
 
 def create_demo():
-    # Initialize the assistant
     assistant = HealthAssistant()
     
-    # Create the Gradio interface
-    with gr.Blocks(css="footer {visibility: hidden}") as demo:
-        gr.Markdown(f"# {APP_TITLE}")
-        gr.Markdown("""
-        Asisten digital ini dirancang untuk membantu Anda berkonsultasi tentang kesehatan wanita.
-        
-        _Catatan: Informasi yang diberikan bersifat umum. Selalu konsultasikan dengan tenaga kesehatan untuk saran yang lebih spesifik._
-        """)
+    # Define the interface
+    with gr.Blocks(title=TITLE) as demo:
+        gr.Markdown(DESCRIPTION)
         
         chatbot = gr.Chatbot(
-            value=[[None, INITIAL_MESSAGE]],
-            height=400
+            label="Chat History",
+            height=600,
+            show_copy_button=True,
         )
         
         with gr.Row():
             msg = gr.Textbox(
-                placeholder="Ketik pertanyaan Anda di sini...",
-                show_label=False,
+                label="Ketik pertanyaan Anda di sini...",
+                placeholder="Contoh: Apa itu PCOS?",
                 scale=9
             )
             submit = gr.Button("Kirim", scale=1)
-            clear = gr.Button("🗑️ Hapus Riwayat", scale=1)
+            
+        clear = gr.Button("🗑️ Hapus Riwayat Chat")
         
         # Set up event handlers
         submit_click = submit.click(
             assistant.respond,
             inputs=[msg, chatbot],
             outputs=[chatbot],
-            queue=True
+            show_progress="full"
         )
-        submit_click.then(lambda: "", None, msg)
+        submit_click.then(lambda: "", None, msg)  # Clear input after sending
         
         msg.submit(
             assistant.respond,
             inputs=[msg, chatbot],
             outputs=[chatbot],
-            queue=True
-        ).then(lambda: "", None, msg)
+            show_progress="full"
+        ).then(lambda: "", None, msg)  # Clear input after sending
         
         clear.click(
             assistant.clear_history,
             outputs=[chatbot],
-            queue=False
+            show_progress=True
         )
+        
+        # Add some CSS styling
+        gr.Markdown("""
+        <style>
+        .gradio-container {
+            max-width: 1200px !important;
+            margin: auto;
+        }
+        </style>
+        """)
     
     return demo
 
-# Create and launch the demo
-demo = create_demo()
+if __name__ == "__main__":
+    demo = create_demo()
+    demo.launch(
+        share=True,
+        server_name="0.0.0.0",
+        server_port=7860,
+        enable_queue=True
+    )