Update app.py

app.py

@@ -1,21 +1,177 @@
 import gradio as gr
 import os
 import spaces
-from transformers import GemmaTokenizer, AutoModelForCausalLM
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
+from langchain_community.vectorstores.faiss import FAISS
+from langchain_huggingface import HuggingFaceEmbeddings
 
-# Set an environment variable
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
+class RAGChatbot:
+    def __init__(self):
+        self.init_models()
+        self.init_vector_store()
+        
+        self.background_prompt = '''
+        As an AI legal assistant, you are a highly trained expert in U.S. and Canadian law. Your purpose is to provide accurate, comprehensive, and professional legal information to assist users with a wide range of legal questions. When answering questions, you should actively ask questions to obtain more information, analyze from different perspectives, and explain your reasoning process to the user. 
+
+        In addition to providing general legal advice and analysis, you are also capable of assisting clients with drafting and reviewing standardized contracts and legal documents. However, your primary role is still to provide personalized legal guidance through interactive conversations with clients.
+
+        Please adhere to the following guidelines:
+
+        1. Clarify the question:
+
+        - Ask questions to clarify the user's specific situation and needs to provide the most relevant and targeted advice. 
+        - However, if the user has already provided sufficient background information, avoid excessively pressing for details. Focus on understanding the core of the issue, rather than unnecessary minutiae.
+
+        2. Gather information:  
+
+        - Identify the key information needed to answer the question and proactively ask the user for these details.
+        - When gathering information, be sure to identify which details are directly relevant to the legal analysis of the case. For information that is not relevant, you don't need to ask too many follow-up questions.
+        - If the user indicates that they have provided all relevant information, accept this and do not continue to demand more details.
+
+        3. Multi-perspective analysis:
+
+        - Evaluate legal issues from different viewpoints, considering various possible interpretations and applications.
+        - Present arguments supporting and opposing specific perspectives to comprehensively clarify complex issues.
+        - In your analysis, strive to balance comprehensiveness and conciseness. Provide thorough analysis, but also ensure that the user can easily understand and absorb the information.
+
+        4. Explain reasoning:
+
+        - Explain the main legal principles, regulations, and case law you consider when analyzing the issue.
+        - Clarify how you apply legal knowledge to the user's specific situation and the logic behind your conclusions.
+        - When explaining your reasoning, use clear and concise language, avoiding excessive length or repetition.
+
+        5. Interactive dialogue:
+
+        - Encourage the user to participate in the discussion, ask follow-up questions, and share their thoughts and concerns.
+        - Dynamically adjust your analysis and recommendations based on new information obtained in the conversation.
+        - In your interactions, be attentive to the user's needs and concerns. If they express satisfaction or indicate that they don't require more information, respect their wishes.
 
+        6. Professional advice:  
 
+        - Provide clear, actionable legal advice, but also emphasize the necessity of consulting a professional lawyer before making a final decision.
+        - If clients wish to speak with a lawyer, you can introduce our team (WealthWizards), which consists of lawyers with different specializations and orientations.
+        - When providing advice, use language that is easy to understand and communicate with a tone of empathy and care. Let them feel that you understand their situation and sincerely want to help them.
 
-MODEL_NAME_OR_PATH = 'StevenChen16/llama3-8b-Lawyer'
-# MODEL_NAME_OR_PATH = 'nvidia/Llama3-ChatQA-1.5-8B'
+        7. Assistance with standardized contracts and legal documents:
 
+        - When clients request assistance with drafting or reviewing standardized contracts and legal documents, provide guidance and support to the best of your abilities.
+        - Analyze the client's needs and requirements, and offer suggestions on appropriate contract templates or clauses to include.
+        - Review drafted documents for potential legal issues or areas that may need improvement, and provide constructive feedback.
+        - However, always remind clients that while you can assist with drafting and review, final documents should still be reviewed and approved by a licensed attorney.
 
+        Please remember that your role is to provide general legal information and analysis, but also to actively guide and interact with the user during the conversation in a personalized and professional manner. If you feel that necessary information is missing to provide targeted analysis and advice, take the initiative to ask until you believe you have sufficient details. However, also be mindful to avoid over-inquiring or disregarding the user's needs and concerns. 
 
+        When assisting with standardized contracts and documents, aim to provide value-added services while still maintaining the importance of attorney review. Your contract assistance should be a supplement to, not a replacement for, the interactive legal guidance that is your primary function.
 
+        Now, please guide me step by step to describe the legal issues I am facing, according to the above requirements.
+        '''
+
+    @spaces.GPU
+    def init_models(self):
+        """Initialize the LLM and embedding models"""
+        # LLM initialization
+        self.llm_model_name = 'StevenChen16/llama3-8b-Lawyer'
+        self.tokenizer = AutoTokenizer.from_pretrained(self.llm_model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.llm_model_name, 
+            device_map="auto"
+        )
+        self.terminators = [
+            self.tokenizer.eos_token_id,
+            self.tokenizer.convert_tokens_to_ids("<|eot_id|>")
+        ]
+        
+        # Embedding model initialization
+        self.embeddings = HuggingFaceEmbeddings(
+            model_name='intfloat/multilingual-e5-large-instruct',
+            model_kwargs={'trust_remote_code': True}
+        )
+
+    def init_vector_store(self):
+        """Load vector store from HuggingFace Hub"""
+        try:
+            self.vector_store = FAISS.from_huggingface_hub(
+                repo_id="StevenChen16/laws.faiss",
+                embeddings=self.embeddings,
+                allow_dangerous_deserialization=True
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to load vector store from HuggingFace Hub: {str(e)}")
+
+    def get_relevant_context(self, query, k=4):
+        """Retrieve relevant context from vector store"""
+        retriever = self.vector_store.as_retriever(
+            search_type="similarity_score_threshold",
+            search_kwargs={
+                "score_threshold": 0.7,
+                "k": k
+            }
+        )
+        docs = retriever.invoke(query)
+        return "\n".join(doc.page_content for doc in docs)
+
+    @spaces.GPU(duration=120)
+    def generate_response(self, message, history, temperature=0.6, max_new_tokens=4096):
+        """Generate streaming response with RAG context"""
+        # Get relevant context
+        context = self.get_relevant_context(message)
+        
+        # Build conversation history
+        conversation = []
+        for user, assistant in history:
+            conversation.extend([
+                {"role": "user", "content": user},
+                {"role": "assistant", "content": assistant}
+            ])
+            
+        # Add context and background prompt to message
+        if context:
+            enhanced_message = f"""Based on the following legal context:
+            {context}
+            
+            {self.background_prompt}
+            
+            User question: {message}"""
+        else:
+            enhanced_message = message + self.background_prompt
+            
+        conversation.append({"role": "user", "content": enhanced_message})
+
+        # Generate response
+        input_ids = self.tokenizer.apply_chat_template(
+            conversation,
+            return_tensors="pt"
+        ).to(self.model.device)
+        
+        streamer = TextIteratorStreamer(
+            self.tokenizer,
+            timeout=10.0,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+
+        generate_kwargs = dict(
+            input_ids=input_ids,
+            streamer=streamer,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=temperature,
+            eos_token_id=self.terminators,
+        )
+        
+        if temperature == 0:
+            generate_kwargs['do_sample'] = False
+            
+        t = Thread(target=self.model.generate, kwargs=generate_kwargs)
+        t.start()
+
+        outputs = []
+        for text in streamer:
+            outputs.append(text)
+            yield "".join(outputs)
+
+# Gradio interface constants
 DESCRIPTION = '''
 <div style="display: flex; align-items: center; justify-content: center; text-align: center;">
     <a href="https://wealthwizards.org/" target="_blank">
@@ -43,7 +199,6 @@ PLACEHOLDER = """
 </div>
 """
 
-
 css = """
 h1 {
   text-align: center;
@@ -57,129 +212,16 @@ h1 {
 }
 """
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_OR_PATH)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME_OR_PATH, device_map="auto")  # to("cuda:0") 
-terminators = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|eot_id|>")
-]
-
-
-background_prompt = '''
-As an AI legal assistant, you are a highly trained expert in U.S. and Canadian law. Your purpose is to provide accurate, comprehensive, and professional legal information to assist users with a wide range of legal questions. When answering questions, you should actively ask questions to obtain more information, analyze from different perspectives, and explain your reasoning process to the user. 
-
-In addition to providing general legal advice and analysis, you are also capable of assisting clients with drafting and reviewing standardized contracts and legal documents. However, your primary role is still to provide personalized legal guidance through interactive conversations with clients.
-
-Please adhere to the following guidelines:
-
-1. Clarify the question:
-
-- Ask questions to clarify the user's specific situation and needs to provide the most relevant and targeted advice. 
-- However, if the user has already provided sufficient background information, avoid excessively pressing for details. Focus on understanding the core of the issue, rather than unnecessary minutiae.
-
-2. Gather information:  
+# Initialize chatbot
+chatbot = RAGChatbot()
 
-- Identify the key information needed to answer the question and proactively ask the user for these details.
-- When gathering information, be sure to identify which details are directly relevant to the legal analysis of the case. For information that is not relevant, you don't need to ask too many follow-up questions.
-- If the user indicates that they have provided all relevant information, accept this and do not continue to demand more details.
-
-3. Multi-perspective analysis:
-
-- Evaluate legal issues from different viewpoints, considering various possible interpretations and applications.
-- Present arguments supporting and opposing specific perspectives to comprehensively clarify complex issues.
-- In your analysis, strive to balance comprehensiveness and conciseness. Provide thorough analysis, but also ensure that the user can easily understand and absorb the information.
-
-4. Explain reasoning:
-
-- Explain the main legal principles, regulations, and case law you consider when analyzing the issue.
-- Clarify how you apply legal knowledge to the user's specific situation and the logic behind your conclusions.
-- When explaining your reasoning, use clear and concise language, avoiding excessive length or repetition.
-
-5. Interactive dialogue:
-
-- Encourage the user to participate in the discussion, ask follow-up questions, and share their thoughts and concerns.
-- Dynamically adjust your analysis and recommendations based on new information obtained in the conversation.
-- In your interactions, be attentive to the user's needs and concerns. If they express satisfaction or indicate that they don't require more information, respect their wishes.
-
-6. Professional advice:  
-
-- Provide clear, actionable legal advice, but also emphasize the necessity of consulting a professional lawyer before making a final decision.
-- If clients wish to speak with a lawyer, you can introduce our team (WealthWizards), which consists of lawyers with different specializations and orientations.
-- When providing advice, use language that is easy to understand and communicate with a tone of empathy and care. Let them feel that you understand their situation and sincerely want to help them.
-
-7. Assistance with standardized contracts and legal documents:
-
-- When clients request assistance with drafting or reviewing standardized contracts and legal documents, provide guidance and support to the best of your abilities.
-- Analyze the client's needs and requirements, and offer suggestions on appropriate contract templates or clauses to include.
-- Review drafted documents for potential legal issues or areas that may need improvement, and provide constructive feedback.
-- However, always remind clients that while you can assist with drafting and review, final documents should still be reviewed and approved by a licensed attorney.
-
-Please remember that your role is to provide general legal information and analysis, but also to actively guide and interact with the user during the conversation in a personalized and professional manner. If you feel that necessary information is missing to provide targeted analysis and advice, take the initiative to ask until you believe you have sufficient details. However, also be mindful to avoid over-inquiring or disregarding the user's needs and concerns. 
-
-When assisting with standardized contracts and documents, aim to provide value-added services while still maintaining the importance of attorney review. Your contract assistance should be a supplement to, not a replacement for, the interactive legal guidance that is your primary function.
-
-Now, please guide me step by step to describe the legal issues I am facing, according to the above requirements.
-'''
-
-@spaces.GPU(duration=120)
-def chat_llama3_8b(message: str, 
-              history: list, 
-              temperature=0.6, 
-              max_new_tokens=4096
-             ) -> str:
-    """
-    Generate a streaming response using the llama3-8b model.
-    Args:
-        message (str): The input message.
-        history (list): The conversation history used by ChatInterface.
-        temperature (float): The temperature for generating the response.
-        max_new_tokens (int): The maximum number of new tokens to generate.
-    Returns:
-        str: The generated response.
-    """
-    conversation = []
-    for user, assistant in history:
-        # content = background_prompt + user
-        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
-    message = message + background_prompt
-    conversation.append({"role": "user", "content": message})
-
-    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
-    
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-
-    generate_kwargs = dict(
-        input_ids= input_ids,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        temperature=temperature,
-        eos_token_id=terminators,
-    )
-    # This will enforce greedy generation (do_sample=False) when the temperature is passed 0, avoiding the crash.             
-    if temperature == 0:
-        generate_kwargs['do_sample'] = False
-        
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        #print(outputs)
-        yield "".join(outputs)
-        
-
-# Gradio block
-chatbot=gr.Chatbot(height=600, placeholder=PLACEHOLDER, label='Gradio ChatInterface')
-
-with gr.Blocks(fill_height=True, css=css) as demo:
-    
+# Gradio interface
+with gr.Blocks(css=css) as demo:
     gr.Markdown(DESCRIPTION)
+    
     gr.ChatInterface(
-        fn=chat_llama3_8b,
-        chatbot=chatbot,
+        fn=chatbot.generate_response,
+        chatbot=gr.Chatbot(height=600, placeholder=PLACEHOLDER, label='AI Legal Assistant'),
         fill_height=True,
         examples=[
             ['What are the key differences between a sole proprietorship and a partnership?'],
@@ -189,9 +231,9 @@ with gr.Blocks(fill_height=True, css=css) as demo:
             ['How can I protect my intellectual property when sharing my idea with potential investors?']
         ],
         cache_examples=False,
-                     )
+    )
     
     gr.Markdown(LICENSE)
-    
+
 if __name__ == "__main__":
     demo.launch()