Merge branch 'main' of https://github.com/arc53/DocsGPT into main

scripts/parser/file/rst_parser.py

@@ -24,6 +24,8 @@ class RstParser(BaseParser):
         remove_hyperlinks: bool = True,
         remove_images: bool = True,
         remove_table_excess: bool = True,
+        remove_interpreters: bool = True,
+        remove_directives: bool = True,
         remove_whitespaces_excess: bool = True,
         #Be carefull with remove_characters_excess, might cause data loss
         remove_characters_excess: bool = True,
@@ -34,6 +36,8 @@ class RstParser(BaseParser):
         self._remove_hyperlinks = remove_hyperlinks
         self._remove_images = remove_images
         self._remove_table_excess = remove_table_excess
+        self._remove_interpreters = remove_interpreters
+        self._remove_directives = remove_directives
         self._remove_whitespaces_excess = remove_whitespaces_excess
         self._remove_characters_excess = remove_characters_excess
 
@@ -95,6 +99,18 @@ class RstParser(BaseParser):
         content = re.sub(pattern, r"\1", content)
         return content
 
+    def remove_directives(self, content: str) -> str:
+        """Removes reStructuredText Directives"""
+        pattern = r"`\.\.([^:]+)::"
+        content = re.sub(pattern, "", content)
+        return content
+
+    def remove_interpreters(self, content: str) -> str:
+        """Removes reStructuredText Interpreted Text Roles"""
+        pattern = r":(\w+):"
+        content = re.sub(pattern, "", content)
+        return content
+
     def remove_table_excess(self, content: str) -> str:
         """Pattern to remove grid table separators"""
         pattern = r"^\+[-]+\+[-]+\+$"
@@ -129,6 +145,10 @@ class RstParser(BaseParser):
             content = self.remove_images(content)
         if self._remove_table_excess:
             content = self.remove_table_excess(content)
+        if self._remove_directives:
+            content = self.remove_directives(content)
+        if self._remove_interpreters:
+            content = self.remove_interpreters(content)
         rst_tups = self.rst_to_tups(content)
         if self._remove_whitespaces_excess:
             rst_tups = self.remove_whitespaces_excess(rst_tups)

scripts/parser/open_ai_func.py

@@ -1,9 +1,17 @@
+import os
 import faiss
 import pickle
 import tiktoken
 from langchain.vectorstores import FAISS
 from langchain.embeddings import OpenAIEmbeddings
 
+#from langchain.embeddings import HuggingFaceEmbeddings
+#from langchain.embeddings import HuggingFaceInstructEmbeddings
+#from langchain.embeddings import CohereEmbeddings
+
+from retry import retry
+
+
 
 def num_tokens_from_string(string: str, encoding_name: str) -> int:
 # Function to convert string to tokens and estimate user cost.
@@ -12,16 +20,44 @@ def num_tokens_from_string(string: str, encoding_name: str) -> int:
     total_price = ((num_tokens/1000) * 0.0004)
     return num_tokens, total_price
 
-def call_openai_api(docs):
+@retry(tries=10, delay=60)
+def store_add_texts_with_retry(store, i):
+    store.add_texts([i.page_content], metadatas=[i.metadata])
+
+def call_openai_api(docs, folder_name):
 # Function to create a vector store from the documents and save it to disk.
-    store = FAISS.from_documents(docs, OpenAIEmbeddings())
-    faiss.write_index(store.index, "docs.index")
-    store.index = None
 
-    with open("faiss_store.pkl", "wb") as f:
-        pickle.dump(store, f)
+    # create output folder if it doesn't exist
+    if not os.path.exists(f"outputs/{folder_name}"):
+        os.makedirs(f"outputs/{folder_name}")
 
-def get_user_permission(docs):
+    from tqdm import tqdm
+    docs_test = [docs[0]]
+    # remove the first element from docs
+    docs.pop(0)
+    # cut first n docs if you want to restart
+    #docs = docs[:n]
+    c1 = 0
+    store = FAISS.from_documents(docs_test, OpenAIEmbeddings())
+
+    # Uncomment for MPNet embeddings
+    # model_name = "sentence-transformers/all-mpnet-base-v2"
+    # hf = HuggingFaceEmbeddings(model_name=model_name)
+    # store = FAISS.from_documents(docs_test, hf)
+    for i in tqdm(docs, desc="Embedding 🦖", unit="docs", total=len(docs), bar_format='{l_bar}{bar}| Time Left: {remaining}'):
+        try:
+            store_add_texts_with_retry(store, i)
+        except Exception as e:
+            print(e)
+            print("Error on ", i)
+            print("Saving progress")
+            print(f"stopped at {c1} out of {len(docs)}")
+            store.save_local(f"outputs/{folder_name}")
+            break
+        c1 += 1
+    store.save_local(f"outputs/{folder_name}")
+
+def get_user_permission(docs, folder_name):
 # Function to ask user permission to call the OpenAI api and spend their OpenAI funds.
     # Here we convert the docs list to a string and calculate the number of OpenAI tokens the string represents.
     #docs_content = (" ".join(docs))
@@ -37,8 +73,8 @@ def get_user_permission(docs):
     #Here we check for user permission before calling the API.
     user_input = input("Price Okay? (Y/N) \n").lower()
     if user_input == "y":
-        call_openai_api(docs)
+        call_openai_api(docs, folder_name)
     elif user_input == "":
-        call_openai_api(docs)
+        call_openai_api(docs, folder_name)
     else:
-        print("The API was not called. No money was spent.")
+        print("The API was not called. No money was spent.")