app.py

app.py

@@ -34,7 +34,7 @@ class BSIChatbot:
         self.docs = docs_path
         self.rerank_model_path = model_paths['rerank_model_path']
 
-    #@spaces.GPU
+    @spaces.GPU
     def initialize_embedding_model(self, rebuild_embeddings: bool):
         raw_knowledge_base = []
 
@@ -80,10 +80,12 @@ class BSIChatbot:
             # Load existing vector store
             self.vectorstore = FAISS.load_local(os.path.join(self.docs, "_embeddings"), self.embedding_model)
 
+    @spaces.GPU
     def retrieve_similar_embedding(self, query: str):
         query = f"Instruct: Given a search query, retrieve the relevant passages that answer the query\nQuery:{query}"
         return self.vectorstore.similarity_search(query=query, k=20)
 
+    @spaces.GPU
     def initialize_llm(self):
         bnb_config = BitsAndBytesConfig(load_in_8bit=True)
         llm = AutoModelForCausalLM.from_pretrained(self.llm_path, quantization_config=bnb_config)
@@ -102,7 +104,7 @@ class BSIChatbot:
             max_new_tokens=500,
         )
 
-
+    @spaces.GPU
     def rag_prompt(self, query: str, rerank: bool, history: List[Dict]):
         retrieved_chunks = self.retrieve_similar_embedding(query)
         retrieved_texts = [f"{chunk.metadata['source']}:\n{chunk.page_content}" for chunk in retrieved_chunks]
@@ -125,7 +127,7 @@ class BSIChatbot:
 
         return self.streamer
 
-    @spaces.GPU
+
     def launch_interface(self):
         with gr.Blocks() as demo:
             chatbot = gr.Chatbot(type="messages")