EddyGiusepe's picture
Primeira tentativa
2d2012b
Raw
History Blame Contribute Delete
1.5 kB
from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
from langchain.document_loaders import UnstructuredFileLoader, DirectoryLoader
from langchain.vectorstores.faiss import FAISS
from langchain.embeddings import OpenAIEmbeddings
import pickle
import os
# loader = UnstructuredFileLoader("state_of_the_union.txt")
def embed_doc():
#check data folder is not empty
if len(os.listdir("data")) > 0:
loader = DirectoryLoader('data', glob="**/*.*")
raw_documents = loader.load()
print(len(raw_documents))
# Split text
text_splitter = RecursiveCharacterTextSplitter(
# Set a really small chunk size, just to show.
chunk_size = 1000,
chunk_overlap = 0,
length_function = len,
)
print("111")
documents = text_splitter.split_documents(raw_documents)
# Load Data to vectorstore
embeddings = OpenAIEmbeddings()
print("222")
vectorstore = FAISS.from_documents(documents, embeddings)
print("333")
# Save vectorstore
# check if vectorstore.pkl exists
with open("vectorstore.pkl", "wb") as f:
pickle.dump(vectorstore, f)
# check if vectorstore.pkl exists
if os.path.exists("vectorstore.pkl"):
with open("vectorstore.pkl", 'rb') as f:
docsearch = pickle.load(f)
# query = input("Enter your query: ")
# docs = docsearch.similarity_search(query)
# print(docs[0])