From 9465ccf6e5977deff939bb26c852df7849a54b93 Mon Sep 17 00:00:00 2001 From: rishiraj Date: Mon, 15 Apr 2024 17:42:45 +0530 Subject: [PATCH] add to_df() feature --- requirements.txt | 3 ++- spanking/main.py | 8 ++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 3edcb95..c31ee6c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ jax==0.4.26 -sentence-transformers==2.6.1 \ No newline at end of file +sentence-transformers==2.6.1 +pandas==2.0.3 \ No newline at end of file diff --git a/spanking/main.py b/spanking/main.py index 98cfb3b..7b5878c 100644 --- a/spanking/main.py +++ b/spanking/main.py @@ -2,6 +2,7 @@ import jax.numpy as jnp import pickle from sentence_transformers import SentenceTransformer +import pandas as pd class VectorDB: def __init__(self, model_name='BAAI/bge-base-en-v1.5'): @@ -46,6 +47,13 @@ def save(self, file_path): def load(file_path): with open(file_path, 'rb') as file: return pickle.load(file) + + def to_df(self): + data = { + 'text': self.texts, + 'embedding': [embedding.tolist() for embedding in self.embeddings] + } + return pd.DataFrame(data) def __len__(self): return len(self.texts)