RustX/CSV-ChatBot

1import os2import pickle3import tempfile4from langchain.document_loaders.csv_loader import CSVLoader5from langchain.vectorstores import FAISS6from langchain.embeddings.openai import OpenAIEmbeddings7 8 9class Embedder:10    def __init__(self):11        self.PATH = "embeddings"12        self.createEmbeddingsDir()13 14    def createEmbeddingsDir(self):15        """16        Creates a directory to store the embeddings vectors17        """18        if not os.path.exists(self.PATH):19            os.mkdir(self.PATH)20 21    def storeDocEmbeds(self, file, filename):22        """23        Stores document embeddings using Langchain and FAISS24        """25        # Write the uploaded file to a temporary file26        with tempfile.NamedTemporaryFile(mode="wb", delete=False) as tmp_file:27            tmp_file.write(file)28            tmp_file_path = tmp_file.name29 30        # Load the data from the file using Langchain31        loader = CSVLoader(file_path=tmp_file_path, encoding="utf-8")32        data = loader.load_and_split()33 34        # Create an embeddings object using Langchain35        embeddings = OpenAIEmbeddings()36 37        # Store the embeddings vectors using FAISS38        vectors = FAISS.from_documents(data, embeddings)39        os.remove(tmp_file_path)40 41        # Save the vectors to a pickle file42        with open(f"{self.PATH}/{filename}.pkl", "wb") as f:43            pickle.dump(vectors, f)44 45    def getDocEmbeds(self, file, filename):46        """47        Retrieves document embeddings48        """49        # Check if embeddings vectors have already been stored in a pickle file50        if not os.path.isfile(f"{self.PATH}/{filename}.pkl"):51            # If not, store the vectors using the storeDocEmbeds function52            self.storeDocEmbeds(file, filename)53 54        # Load the vectors from the pickle file55        with open(f"{self.PATH}/{filename}.pkl", "rb") as f:56            vectors = pickle.load(f)57 58        return vectors