Spaces:
Runtime error
Runtime error
File size: 2,012 Bytes
c0eb7b0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import os
import pickle
import tempfile
from langchain.document_loaders import PyPDFLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
class Embedder:
def __init__(self):
self.PATH = "embeddings"
self.createEmbeddingsDir()
def createEmbeddingsDir(self):
"""
Creates a directory to store the embeddings vectors
"""
if not os.path.exists(self.PATH):
os.mkdir(self.PATH)
def storeDocEmbeds(self, file, filename):
"""
Stores document embeddings using Langchain and FAISS
"""
# Write the uploaded file to a temporary file
with tempfile.NamedTemporaryFile(mode="wb", delete=False) as tmp_file:
tmp_file.write(file)
tmp_file_path = tmp_file.name
# Load the data from the file using Langchain
loader = PyPDFLoader(file_path=tmp_file_path)
data = loader.load_and_split()
print(f"Loaded {len(data)} documents from {tmp_file_path}")
# Create an embeddings object using Langchain
embeddings = OpenAIEmbeddings(allowed_special={'<|endofprompt|>'})
# Store the embeddings vectors using FAISS
vectors = FAISS.from_documents(data, embeddings)
os.remove(tmp_file_path)
# Save the vectors to a pickle file
with open(f"{self.PATH}/{filename}.pkl", "wb") as f:
pickle.dump(vectors, f)
def getDocEmbeds(self, file, filename):
"""
Retrieves document embeddings
"""
# Check if embeddings vectors have already been stored in a pickle file
pkl_file = f"{self.PATH}/{filename}.pkl"
if not os.path.isfile(pkl_file):
# If not, store the vectors using the storeDocEmbeds function
self.storeDocEmbeds(file, filename)
# Load the vectors from the pickle file
with open(pkl_file, "rb") as f:
vectors = pickle.load(f)
return vectors |