Update app.py

app.py

@@ -68,10 +68,41 @@ class RAGChatbot:
         Now, please guide me step by step to describe the legal issues I am facing, according to the above requirements.
         '''
 
+    import gradio as gr
+import os
+import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+from langchain_community.vectorstores.faiss import FAISS
+from langchain_huggingface import HuggingFaceEmbeddings
+from huggingface_hub import snapshot_download
+
+class RAGChatbot:
+    def __init__(self):
+        # First initialize models to create embeddings
+        self.init_models()
+        # Then initialize vector store which uses embeddings
+        self.init_vector_store()
+        
+        self.background_prompt = '''
+        As an AI legal assistant, you are a highly trained expert in U.S. and Canadian law...
+        [rest of your existing background prompt]
+        '''
+
     @spaces.GPU
     def init_models(self):
         """Initialize the LLM and embedding models"""
+        print("Initializing models...")
+        
+        # Embedding model initialization first
+        print("Loading embedding model...")
+        self.embeddings = HuggingFaceEmbeddings(
+            model_name='intfloat/multilingual-e5-large-instruct',
+            model_kwargs={'trust_remote_code': True}
+        )
+        
         # LLM initialization
+        print("Loading LLM model...")
         self.llm_model_name = 'StevenChen16/llama3-8b-Lawyer'
         self.tokenizer = AutoTokenizer.from_pretrained(self.llm_model_name)
         self.model = AutoModelForCausalLM.from_pretrained(
@@ -82,28 +113,27 @@ class RAGChatbot:
             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}
-        )
+        print("Models initialized successfully")
 
     def init_vector_store(self):
         """Load vector store from HuggingFace Hub"""
         try:
+            print("Downloading vector store from HuggingFace Hub...")
             # Download FAISS files from HuggingFace Hub
             repo_path = snapshot_download(
                 repo_id="StevenChen16/laws.faiss",
                 repo_type="model"
             )
             
+            print("Loading vector store...")
             # Load the vector store from downloaded files
             self.vector_store = FAISS.load_local(
-                repo_path,
-                self.embeddings,
+                folder_path=repo_path,  # Specify the parameter name explicitly
+                embeddings=self.embeddings,
                 allow_dangerous_deserialization=True
             )
+            print("Vector store loaded successfully")
+            
         except Exception as e:
             raise RuntimeError(f"Failed to load vector store from HuggingFace Hub: {str(e)}")