Spaces:
Sleeping
Sleeping
import sys | |
import os | |
import pinecone | |
from langchain.text_splitter import RecursiveCharacterTextSplitter | |
from langchain.document_loaders import DirectoryLoader | |
from langchain.embeddings import OpenAIEmbeddings | |
from langchain.vectorstores import Pinecone | |
PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY") | |
PINECONE_API_ENV = os.environ.get("PINECONE_API_ENV") | |
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY") | |
PINECONE_INDEX_NAME = os.environ.get("PINECONE_INDEX_NAME") | |
def load_documents(path_to_files): | |
# Uses UnstructuredLoader under the hood | |
loader = DirectoryLoader(path=path_to_files, glob="*.json") | |
raw_documents = loader.load() | |
text_splitter = RecursiveCharacterTextSplitter() | |
documents = text_splitter.split_documents(raw_documents) | |
return documents | |
def send_docs_to_pinecone(documents): | |
embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY) | |
pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_API_ENV) | |
if PINECONE_INDEX_NAME in pinecone.list_indexes(): | |
print( | |
f"Index {PINECONE_INDEX_NAME} already exists, deleting and recreating to avoid duplicates" | |
) | |
pinecone.delete_index(name=PINECONE_INDEX_NAME) | |
pinecone.create_index(name=PINECONE_INDEX_NAME, dimension=1536) | |
Pinecone.from_documents(documents, embeddings, index_name=PINECONE_INDEX_NAME) | |
if __name__ == "__main__": | |
path_to_files = sys.argv[1] | |
print(f"Grabbing json files from {path_to_files}") | |
docs = load_documents(path_to_files) | |
print(f"Found {len(docs)}, sending to pinecone") | |
send_docs_to_pinecone(docs) | |