LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2025-11-29 16:43:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add support for directory list

Browse Source 
example: `python ingest.py --dir inputs1 --dir another --dir ../inputs`,
the outputs will be in `outputs/input_folder_name/`
This commit is contained in:
冯不游
2023-02-15 02:30:39 +08:00
parent 5883ce2685
commit b83589a308
2 changed files with 56 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
scripts/parser/open_ai_func.py
View File

@@ -1,3 +1,4 @@
import os
import faiss
import pickle
import tiktoken
@@ -12,8 +13,13 @@ 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):
def call_openai_api(docs, folder_name):
# Function to create a vector store from the documents and save it to disk.
# create output folder if it doesn't exist
if not os.path.exists(f"outputs/{folder_name}"):
os.makedirs(f"outputs/{folder_name}")
from tqdm import tqdm
docs_test = [docs[0]]
# remove the first element from docs
@@ -31,25 +37,23 @@ def call_openai_api(docs):
print("Error on ", i)
print("Saving progress")
print(f"stopped at {c1} out of {len(docs)}")
faiss.write_index(store.index, "docs.index")
faiss.write_index(store.index, f"outputs/{folder_name}/docs.index")
store_index_bak = store.index
store.index = None
with open("faiss_store.pkl", "wb") as f:
with open(f"outputs/{folder_name}/faiss_store.pkl", "wb") as f:
pickle.dump(store, f)
print("Sleeping for 60 seconds and trying again")
time.sleep(60)
faiss.write_index(store_index_bak, "docs.index")
store.index = store_index_bak
store.add_texts([i.page_content], metadatas=[i.metadata])
c1 += 1
faiss.write_index(store.index, "docs.index")
faiss.write_index(store.index, f"outputs/{folder_name}/docs.index")
store.index = None
with open("faiss_store.pkl", "wb") as f:
with open(f"outputs/{folder_name}/faiss_store.pkl", "wb") as f:
pickle.dump(store, f)
def get_user_permission(docs):
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))
@@ -65,8 +69,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.")