Update Classes/Helper_Class.py

Classes/Helper_Class.py

@@ -1,57 +1,48 @@
-from typing import List
-from PyPDF2 import PdfReader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_google_genai import GoogleGenerativeAIEmbeddings
-from langchain_community.vectorstores import FAISS
-import os 
-import google.generativeai as genai
-
-os.environ["GOOGLE_API_KEY"] = "AIzaSyBoghqvvnMMS4bA61LjQkkPNdIRetqk438"
-genai.configure(api_key="AIzaSyBoghqvvnMMS4bA61LjQkkPNdIRetqk438")
-
-
-class GenerateFIASSDB:
-    def __init__(self,pdf_docs : List[str], save_loc:str, model_embeddings: str = "models/embedding-001")-> None:
-        self.save_loc = save_loc
-        self.embedding = model_embeddings
-        text = self.get_pdf_text(pdf_docs)
-        text_chunks = self.get_text_chunks(text)
-        self.get_vector_store(text_chunks)
-        pass  #configure gen ai key from config file
-
-    def get_pdf_text(self,pdf_docs : List[str]) -> str:
-        text = ""
-        for pdf in pdf_docs:
-            pdf_reader= PdfReader(pdf)
-            for page in pdf_reader.pages:
-                text+= page.extract_text()
-        return text
-    
-    def get_text_chunks(self, text : str) -> List:
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
-        chunks = text_splitter.split_text(text)
-        return chunks
-    
-    def get_vector_store(self, text_chunks : List) -> None:
-        embeddings = GoogleGenerativeAIEmbeddings(model = self.embedding)
-        vector_store = FAISS.from_texts(text_chunks, embedding=embeddings)
-        vector_store.save_local(self.save_loc)
-
-
-class DB_Retriever:
-    def __init__(self, db_loc : str, model_embeddings : str = "models/embedding-001") -> None:
-        self.db_loc = db_loc
-        self.embeddings = GoogleGenerativeAIEmbeddings(model = model_embeddings)
-        self.db = FAISS.load_local(self.db_loc, self.embeddings,allow_dangerous_deserialization  = True)
-    
-    def retrieve(self, query : str) -> List[str]:
-        # docs = self.db.similarity_search(query)
-        retriver = self.db.as_retriever()
-        # output_docs =  retriver.invoke(query)
-        # return output_docs
-        return retriver
-
-if __name__ =="__main__":
-    res = DB_Retriever("src/faiss_index").retrieve("What is cloud adapter in google connection?")
-    print(len(res))
-    print('\n\n\n\n',res[1])
+from typing import List
+from PyPDF2 import PdfReader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+from langchain_community.vectorstores import FAISS
+import google.generativeai as genai
+
+
+class GenerateFIASSDB:
+    def __init__(self,pdf_docs : List[str], save_loc:str, model_embeddings: str = "models/embedding-001")-> None:
+        self.save_loc = save_loc
+        self.embedding = model_embeddings
+        text = self.get_pdf_text(pdf_docs)
+        text_chunks = self.get_text_chunks(text)
+        self.get_vector_store(text_chunks)
+        pass  #configure gen ai key from config file
+
+    def get_pdf_text(self,pdf_docs : List[str]) -> str:
+        text = ""
+        for pdf in pdf_docs:
+            pdf_reader= PdfReader(pdf)
+            for page in pdf_reader.pages:
+                text+= page.extract_text()
+        return text
+    
+    def get_text_chunks(self, text : str) -> List:
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
+        chunks = text_splitter.split_text(text)
+        return chunks
+    
+    def get_vector_store(self, text_chunks : List) -> None:
+        embeddings = GoogleGenerativeAIEmbeddings(model = self.embedding)
+        vector_store = FAISS.from_texts(text_chunks, embedding=embeddings)
+        vector_store.save_local(self.save_loc)
+
+
+class DB_Retriever:
+    def __init__(self, db_loc : str, model_embeddings : str = "models/embedding-001") -> None:
+        self.db_loc = db_loc
+        self.embeddings = GoogleGenerativeAIEmbeddings(model = model_embeddings)
+        self.db = FAISS.load_local(self.db_loc, self.embeddings,allow_dangerous_deserialization  = True)
+    
+    def retrieve(self, query : str) -> List[str]:
+        # docs = self.db.similarity_search(query)
+        retriver = self.db.as_retriever()
+        # output_docs =  retriver.invoke(query)
+        # return output_docs
+        return retriver