Merge pull request #6 from rishiraj/staging

Staging

Author: rishiraj

requirements.txt

@@ -0,0 +1,2 @@
+jax==0.4.26
+sentence-transformers==2.6.1

setup.py

@@ -26,6 +26,9 @@
 with open("README.md", "r", encoding="utf-8") as fh:
     long_description = fh.read()
 
+with open('requirements.txt') as fh:
+    requirements = fh.read().splitlines()
+
 setuptools.setup(
     name="spanking",
     version=spanking_version,
@@ -45,7 +48,5 @@
     ],
     python_requires=">=3.9",
     entry_points={"console_scripts": ["spanking = spanking.main:main"]},
-    install_requires=[
-        "numpy >= 1.26.4",
-    ],
+    install_requires=requirements,
 )

spanking/main.py

@@ -1,5 +1,6 @@
 import jax
 import jax.numpy as jnp
+import pickle
 from sentence_transformers import SentenceTransformer
 
 class VectorDB:
@@ -37,6 +38,15 @@ def search(self, query, top_k=5):
         top_indices = jnp.argsort(similarities)[-top_k:][::-1]
         return [(self.texts[i], similarities[i]) for i in top_indices]
 
+    def save(self, file_path):
+        with open(file_path, 'wb') as file:
+            pickle.dump(self, file)
+
+    @staticmethod
+    def load(file_path):
+        with open(file_path, 'rb') as file:
+            return pickle.load(file)
+    
     def __len__(self):
         return len(self.texts)