diff --git a/scripts/ingest_rst.py b/scripts/ingest_rst.py index 38881daf..a23cbedc 100644 --- a/scripts/ingest_rst.py +++ b/scripts/ingest_rst.py @@ -5,14 +5,48 @@ from langchain.vectorstores import FAISS from langchain.embeddings import OpenAIEmbeddings import pickle import dotenv +import tiktoken +import sys +def num_tokens_from_string(string: str, encoding_name: str) -> int: +# Function to convert string to tokens and estimate user cost. + encoding = tiktoken.get_encoding(encoding_name) + num_tokens = len(encoding.encode(string)) + total_price = ((num_tokens/1000) * 0.0004) + return num_tokens, total_price + +def call_openai_api(): +# Function to create a vector store from the documents and save it to disk. + store = FAISS.from_texts(docs, OpenAIEmbeddings(), metadatas=metadatas) + faiss.write_index(store.index, "docs.index") + store.index = None + with open("faiss_store.pkl", "wb") as f: + pickle.dump(store, f) + +def get_user_permission(): +# 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)) + tokens, total_price = num_tokens_from_string(string=docs_content, encoding_name="cl100k_base") + # Here we print the number of tokens and the approx user cost with some visually appealing formatting. + print(f"Number of Tokens = {format(tokens, ',d')}") + print(f"Approx Cost = ${format(total_price, ',.2f')}") + #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() + elif user_input == "": + call_openai_api() + else: + print("The API was not called. No money was spent.") + +#Load .env file dotenv.load_dotenv() - # Here we load in the data in the format that Notion exports it in. ps = list(Path("scikit-learn").glob("**/*.rst")) -# parse all child directories +# parse all child directories data = [] sources = [] for p in ps: @@ -30,10 +64,14 @@ for i, d in enumerate(data): docs.extend(splits) metadatas.extend([{"source": sources[i]}] * len(splits)) - -# Here we create a vector store from the documents and save it to disk. -store = FAISS.from_texts(docs, OpenAIEmbeddings(), metadatas=metadatas) -faiss.write_index(store.index, "docs.index") -store.index = None -with open("faiss_store.pkl", "wb") as f: - pickle.dump(store, f) \ No newline at end of file +# Here we check for command line arguments for bot calls. +# If no argument exists or the permission_bypass_flag argument is not '-y', +# user permission is requested to call the API. +if len(sys.argv) > 1: + permission_bypass_flag = sys.argv[1] + if permission_bypass_flag == '-y': + call_openai_api() + else: + get_user_permission() +else: + get_user_permission() \ No newline at end of file diff --git a/scripts/ingest_rst_sphinx.py b/scripts/ingest_rst_sphinx.py index 183ccac9..f2a92638 100644 --- a/scripts/ingest_rst_sphinx.py +++ b/scripts/ingest_rst_sphinx.py @@ -1,6 +1,8 @@ import os import pickle import dotenv +import tiktoken +import sys import faiss import shutil from pathlib import Path @@ -28,6 +30,38 @@ def convert_rst_to_txt(src_dir, dst_dir): f"-C {dst_dir} " sphinx_main(args.split()) +def num_tokens_from_string(string: str, encoding_name: str) -> int: +# Function to convert string to tokens and estimate user cost. + encoding = tiktoken.get_encoding(encoding_name) + num_tokens = len(encoding.encode(string)) + total_price = ((num_tokens/1000) * 0.0004) + return num_tokens, total_price + +def call_openai_api(): +# Function to create a vector store from the documents and save it to disk. + store = FAISS.from_texts(docs, OpenAIEmbeddings(), metadatas=metadatas) + faiss.write_index(store.index, "docs.index") + store.index = None + with open("faiss_store.pkl", "wb") as f: + pickle.dump(store, f) + +def get_user_permission(): +# 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)) + tokens, total_price = num_tokens_from_string(string=docs_content, encoding_name="cl100k_base") + # Here we print the number of tokens and the approx user cost with some visually appealing formatting. + print(f"Number of Tokens = {format(tokens, ',d')}") + print(f"Approx Cost = ${format(total_price, ',.2f')}") + #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() + elif user_input == "": + call_openai_api() + else: + print("The API was not called. No money was spent.") + #Load .env file dotenv.load_dotenv() @@ -58,13 +92,17 @@ for i, d in enumerate(data): docs.extend(splits) metadatas.extend([{"source": sources[i]}] * len(splits)) - -# Here we create a vector store from the documents and save it to disk. -store = FAISS.from_texts(docs, OpenAIEmbeddings(), metadatas=metadatas) -faiss.write_index(store.index, "docs.index") -store.index = None -with open("faiss_store.pkl", "wb") as f: - pickle.dump(store, f) +# Here we check for command line arguments for bot calls. +# If no argument exists or the permission_bypass_flag argument is not '-y', +# user permission is requested to call the API. +if len(sys.argv) > 1: + permission_bypass_flag = sys.argv[1] + if permission_bypass_flag == '-y': + call_openai_api() + else: + get_user_permission() +else: + get_user_permission() # Delete tmp folder # Commented out for now