Update retriever.py

retriever.py

@@ -14,6 +14,38 @@ import os
 curr_dir = os.getcwd()
 db_path = 'chroma_db_v2'
 
+class QuestionRetriever:
+   
+    def load_documents(self,file_name):
+      current_directory = os.getcwd()
+      data_directory = os.path.join(current_directory, "data")
+      file_path = os.path.join(data_directory, file_name)
+      loader = TextLoader(file_path)
+      documents = loader.load()
+      return documents
+    
+    def store_data_in_vector_db(self,documents):
+    #   global db
+      text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0,separator="\n")
+      docs = text_splitter.split_documents(documents)
+      # create the open-source embedding function
+      embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
+      # print(docs)
+      # load it into Chroma
+      db = Chroma.from_documents(docs, embedding_function)
+      return db
+
+    def get_response(self, user_query):
+        db=self.store_data_in_vector_db(documents)
+
+        docs = db.similarity_search(user_query)
+        most_similar_question = docs[0].page_content.split("\n")[0]  # Extract the first question
+        if user_query==most_similar_question:
+          most_similar_question=docs[1].page_content.split("\n")[0]
+
+        print(most_similar_question)
+        return most_similar_question
+        
 def process_pdf_document(file_path, parent_chunk_size=2000, child_chunk_size=500):
     '''
     Process a PDF document and return the documents and text splitters