Merge pull request #12 from rishiraj/staging

Staging

requirements.txt

@@ -1,2 +1,3 @@
 jax==0.4.26
-sentence-transformers==2.6.1
+sentence-transformers==2.6.1
+pandas==2.0.3

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)