Update rag.py

rag.py

@@ -21,10 +21,9 @@ PDF_URL       = "https://arxiv.org/pdf/2303.08774.pdf"
 WEB_URL       = "https://openai.com/research/gpt-4"
 YOUTUBE_URL_1 = "https://www.youtube.com/watch?v=--khbXchTeE"
 YOUTUBE_URL_2 = "https://www.youtube.com/watch?v=hdhZwyf24mE"
-YOUTUBE_URL_3 = "https://www.youtube.com/watch?v=vw-KWfKwvTQ"
 
-YOUTUBE_DIR = "/data/youtube"
-CHROMA_DIR  = "/data/chroma"
+YOUTUBE_DIR = "/data/yt"
+CHROMA_DIR  = "/data/db"
 
 MONGODB_ATLAS_CLUSTER_URI = os.environ["MONGODB_ATLAS_CLUSTER_URI"]
 MONGODB_DB_NAME           = "langchain_db"
@@ -37,45 +36,52 @@ RAG_CHAIN_PROMPT = PromptTemplate(input_variables = ["context", "question"], tem
 client = MongoClient(MONGODB_ATLAS_CLUSTER_URI)
 collection = client[MONGODB_DB_NAME][MONGODB_COLLECTION_NAME]
 
-def document_loading_splitting():
-    # Document loading
+def document_loading():
     docs = []
     
-    # Load PDF
+    # PDF
     loader = PyPDFLoader(PDF_URL)
     docs.extend(loader.load())
     
-    # Load Web
+    # Web
     loader = WebBaseLoader(WEB_URL)
     docs.extend(loader.load())
     
-    # Load YouTube
+    # YouTube
     loader = GenericLoader(YoutubeAudioLoader([YOUTUBE_URL_1,
-                                               YOUTUBE_URL_2,
-                                               YOUTUBE_URL_3], YOUTUBE_DIR), 
+                                               YOUTUBE_URL_2], YOUTUBE_DIR), 
                            OpenAIWhisperParser())
     docs.extend(loader.load())
 
-    # Document splitting
+    returns docs
+
+def document_splitting(config, docs):
     text_splitter = RecursiveCharacterTextSplitter(chunk_overlap = config["chunk_overlap"],
                                                    chunk_size = config["chunk_size"])
-    split_documents = text_splitter.split_documents(docs)
     
-    return split_documents
-
-def document_storage_chroma(documents):
-    Chroma.from_documents(documents = documents, 
+    return text_splitter.split_documents(docs)
+    
+def document_storage_chroma(chunks):
+    Chroma.from_documents(documents = chunks, 
                           embedding = OpenAIEmbeddings(disallowed_special = ()), 
                           persist_directory = CHROMA_DIR)
 
-def document_storage_mongodb(documents):
-    MongoDBAtlasVectorSearch.from_documents(documents = documents,
+def document_storage_mongodb(chunks):
+    MongoDBAtlasVectorSearch.from_documents(documents = chunks,
                                             embedding = OpenAIEmbeddings(disallowed_special = ()),
                                             collection = collection,
                                             index_name = MONGODB_INDEX_NAME)
 
+def rag_batch(config):
+    docs = document_loading()
+    
+    chunks = document_splitting(config, docs)
+    
+    document_storage_chroma(chunks)
+    document_storage_mongodb(chunks)
+
 def document_retrieval_chroma():
-    return Chroma(embedding_function = OpenAIEmbeddings(),
+    return Chroma(embedding_function = OpenAIEmbeddings(disallowed_special = ()),
                   persist_directory = CHROMA_DIR)
 
 def document_retrieval_mongodb():