Spaces:
Sleeping
Sleeping
File size: 1,420 Bytes
2ac805e ed1d17e 2ac805e ed1d17e 2ac805e ed1d17e 2ac805e ed1d17e 2ac805e ed1d17e | 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 | import pandas as pd
from langchain_community.vectorstores import FAISS
from sentence_transformers import SentenceTransformer
from langchain.document_loaders import DataFrameLoader
from langchain.text_splitter import CharacterTextSplitter
class VectorStore:
def __init__(self, persist_dir="vector_db"):
self.persist_dir = persist_dir
self.store = None
# Load the SentenceTransformer model
self.embeddings_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
self.dfs = {}
def add_dataframe(self, df, source_name):
self.dfs[source_name] = df
docs = DataFrameLoader(df).load()
chunks = CharacterTextSplitter(chunk_size=500, chunk_overlap=50).split_documents(docs)
if not self.store:
self.store = FAISS.from_documents(chunks, self.embeddings_model)
else:
self.store.merge_from(FAISS.from_documents(chunks, self.embeddings_model))
def as_retriever(self):
return self.store.as_retriever() if self.store else None
def get_all_dataframes(self):
return self.dfs
def encode_sentences(self, sentences):
# Generate embeddings for the given sentences
return self.embeddings_model.encode(sentences)
def calculate_similarity(self, embeddings):
from sklearn.metrics.pairwise import cosine_similarity
return cosine_similarity(embeddings)
|