Adding implementations for vector stores

.github/tests/lm_tests.py

@@ -8,6 +8,7 @@
 from lotus.cache import CacheConfig, CacheFactory, CacheType
 from lotus.models import LM, SentenceTransformersRM
 from lotus.types import CascadeArgs
+from lotus.vector_store import FaissVS
 
 ################################################################################
 # Setup
@@ -289,7 +290,8 @@ def test_filter_cascade(setup_models):
 def test_join_cascade(setup_models):
     models = setup_models
     rm = SentenceTransformersRM(model="intfloat/e5-base-v2")
-    lotus.settings.configure(lm=models["gpt-4o-mini"], rm=rm)
+    vs = FaissVS() 
+    lotus.settings.configure(lm=models["gpt-4o-mini"], rm=rm, vs=vs)
 
     data1 = {
         "School": [

.github/tests/multimodality_tests.py

@@ -6,6 +6,7 @@
 import lotus
 from lotus.dtype_extensions import ImageArray
 from lotus.models import LM, SentenceTransformersRM
+from lotus.vector_store import FaissVS
 
 ################################################################################
 # Setup
@@ -160,7 +161,8 @@ def test_topk_with_groupby_operation(setup_models, model):
 @pytest.mark.parametrize("model", get_enabled("clip-ViT-B-32"))
 def test_search_operation(setup_models, model):
     rm = setup_models[model]
-    lotus.settings.configure(rm=rm)
+    vs = FaissVS() 
+    lotus.settings.configure(rm=rm, vs=vs)
 
     image_url = [
         "https://img.etsystatic.com/il/4bee20/1469037676/il_340x270.1469037676_iiti.jpg?version=0",
@@ -180,7 +182,8 @@ def test_search_operation(setup_models, model):
 @pytest.mark.parametrize("model", get_enabled("clip-ViT-B-32"))
 def test_sim_join_operation_image_index(setup_models, model):
     rm = setup_models[model]
-    lotus.settings.configure(rm=rm)
+    vs = FaissVS() 
+    lotus.settings.configure(rm=rm, vs=vs)
 
     image_url = [
         "https://img.etsystatic.com/il/4bee20/1469037676/il_340x270.1469037676_iiti.jpg?version=0",
@@ -205,7 +208,8 @@ def test_sim_join_operation_image_index(setup_models, model):
 @pytest.mark.parametrize("model", get_enabled("clip-ViT-B-32"))
 def test_sim_join_operation_text_index(setup_models, model):
     rm = setup_models[model]
-    lotus.settings.configure(rm=rm)
+    vs = FaissVS() 
+    lotus.settings.configure(rm=rm, vs=vs)
 
     image_url = [
         "https://img.etsystatic.com/il/4bee20/1469037676/il_340x270.1469037676_iiti.jpg?version=0",

.github/tests/rm_tests.py

@@ -5,6 +5,7 @@
 
 import lotus
 from lotus.models import CrossEncoderReranker, LiteLLMRM, SentenceTransformersRM
+from lotus.vector_store import ChromaVS, FaissVS, PineconeVS, QdrantVS, WeaviateVS
 
 ################################################################################
 # Setup
@@ -30,6 +31,14 @@
     "text-embedding-3-small": LiteLLMRM,
 }
 
+VECTOR_STORE_TO_CLS = {
+    'local': FaissVS,
+    'weaviate':WeaviateVS,
+    'pinecone': PineconeVS,
+    'chroma': ChromaVS,
+    'qdrant': QdrantVS
+}
+
 
 def get_enabled(*candidate_models: str) -> list[str]:
     return [model for model in candidate_models if model in ENABLED_MODEL_NAMES]
@@ -41,16 +50,28 @@ def setup_models():
 
     for model_name in ENABLED_MODEL_NAMES:
         models[model_name] = MODEL_NAME_TO_CLS[model_name](model=model_name)
+
+
     return models
 
 
+@pytest.fixture(scope='session')
+def setup_vs():
+    vs_model = {}
+
+    for vs in VECTOR_STORE_TO_CLS:
+        vs_model[vs] = VECTOR_STORE_TO_CLS[vs]()
+
+    return vs_model
+
 ################################################################################
 # RM Only Tests
 ################################################################################
 @pytest.mark.parametrize("model", get_enabled("intfloat/e5-small-v2", "text-embedding-3-small"))
 def test_cluster_by(setup_models, model):
     rm = setup_models[model]
-    lotus.settings.configure(rm=rm)
+    vs = FaissVS()
+    lotus.settings.configure(rm=rm, vs=vs)
 
     data = {
         "Course Name": [
@@ -79,7 +100,9 @@ def test_cluster_by(setup_models, model):
 @pytest.mark.parametrize("model", get_enabled("intfloat/e5-small-v2", "text-embedding-3-small"))
 def test_search_rm_only(setup_models, model):
     rm = setup_models[model]
-    lotus.settings.configure(rm=rm)
+    vs = FaissVS()
+
+    lotus.settings.configure(rm=rm, vs=vs)
 
     data = {
         "Course Name": [
@@ -98,7 +121,8 @@ def test_search_rm_only(setup_models, model):
 @pytest.mark.parametrize("model", get_enabled("intfloat/e5-small-v2", "text-embedding-3-small"))
 def test_sim_join(setup_models, model):
     rm = setup_models[model]
-    lotus.settings.configure(rm=rm)
+    vs = FaissVS() 
+    lotus.settings.configure(rm=rm, vs=vs)
 
     data1 = {
         "Course Name": [
@@ -124,7 +148,8 @@ def test_sim_join(setup_models, model):
 )
 def test_dedup(setup_models):
     rm = setup_models["intfloat/e5-small-v2"]
-    lotus.settings.configure(rm=rm)
+    vs = FaissVS() 
+    lotus.settings.configure(rm=rm,vs=vs)
     data = {
         "Text": [
             "Probability and Random Processes",
@@ -142,6 +167,113 @@ def test_dedup(setup_models):
     assert "Probability" in kept[1], kept
 
 
+
+################################################################################
+# VS Only Tests
+################################################################################
+
+
+@pytest.mark.parametrize("vs", VECTOR_STORE_TO_CLS.keys())
+@pytest.mark.parametrize("model", get_enabled("intfloat/e5-small-v2", "text-embedding-3-small"))
+def test_vs_cluster_by(setup_models, setup_vs, vs, model):
+    rm = setup_models[model]
+    my_vs = setup_vs[vs]
+    lotus.settings.configure(rm=rm, vs=my_vs)
+
+    data = {
+        "Course Name": [
+            "Probability and Random Processes",
+            "Cooking",
+            "Food Sciences",
+            "Optimization Methods in Engineering",
+        ]
+    }
+    df = pd.DataFrame(data)
+    df = df.sem_index("Course Name", "indexdir")
+    df = df.sem_cluster_by("Course Name", 2)
+    groups = df.groupby("cluster_id")["Course Name"].apply(set).to_dict()
+    assert len(groups) == 2, groups
+    if "Cooking" in groups[0]:
+        cooking_group = groups[0]
+        probability_group = groups[1]
+    else:
+        cooking_group = groups[1]
+        probability_group = groups[0]
+
+    assert cooking_group == {"Cooking", "Food Sciences"}, groups
+    assert probability_group == {"Probability and Random Processes", "Optimization Methods in Engineering"}, groups
+
+@pytest.mark.parametrize("vs", VECTOR_STORE_TO_CLS.keys())
+@pytest.mark.parametrize("model", get_enabled("intfloat/e5-small-v2", "text-embedding-3-small"))
+def test_vs_search_rm_only(setup_models, setup_vs, vs, model):
+    rm = setup_models[model]
+    my_vs = setup_vs[vs]
+    lotus.settings.configure(rm=rm, vs=my_vs)
+
+    data = {
+        "Course Name": [
+            "Probability and Random Processes",
+            "Cooking",
+            "Food Sciences",
+            "Optimization Methods in Engineering",
+        ]
+    }
+    df = pd.DataFrame(data)
+    df = df.sem_index("Course Name", "secondindexdir")
+    df = df.sem_search("Course Name", "Optimization", K=1)
+    assert df["Course Name"].tolist() == ["Optimization Methods in Engineering"]
+
+@pytest.mark.parametrize("vs", VECTOR_STORE_TO_CLS.keys())
+@pytest.mark.parametrize("model", get_enabled("intfloat/e5-small-v2", "text-embedding-3-small"))
+def test_vs_sim_join(setup_models, setup_vs, vs, model):
+    rm = setup_models[model]
+    my_vs = setup_vs[vs]
+    lotus.settings.configure(rm=rm, vs=my_vs)
+
+    data1 = {
+        "Course Name": [
+            "History of the Atlantic World",
+            "Riemannian Geometry",
+        ]
+    }
+
+    data2 = {"Skill": ["Math", "History"]}
+
+    df1 = pd.DataFrame(data1)
+    df2 = pd.DataFrame(data2).sem_index("Skill", "thirdindexdir")
+    joined_df = df1.sem_sim_join(df2, left_on="Course Name", right_on="Skill", K=1)
+    joined_pairs = set(zip(joined_df["Course Name"], joined_df["Skill"]))
+    expected_pairs = {("History of the Atlantic World", "History"), ("Riemannian Geometry", "Math")}
+    assert joined_pairs == expected_pairs, joined_pairs
+
+
+# TODO: threshold is hardcoded for intfloat/e5-small-v2
+@pytest.mark.skipif(
+    "intfloat/e5-small-v2" not in ENABLED_MODEL_NAMES,
+    reason="Skipping test because intfloat/e5-small-v2 is not enabled",
+)
+@pytest.mark.parametrize("vs", VECTOR_STORE_TO_CLS.keys())
+def test_vs_dedup(setup_models, setup_vs, vs):
+    rm = setup_models["intfloat/e5-small-v2"]
+    my_vs = setup_vs[vs]
+    lotus.settings.configure(rm=rm, vs=my_vs)
+    data = {
+        "Text": [
+            "Probability and Random Processes",
+            "Probability and Markov Chains",
+            "Harry Potter",
+            "Harry James Potter",
+        ]
+    }
+    df = pd.DataFrame(data)
+    df = df.sem_index("Text", "fourthindexdir").sem_dedup("Text", threshold=0.85)
+    kept = df["Text"].tolist()
+    kept.sort()
+    assert len(kept) == 2, kept
+    assert "Harry" in kept[0], kept
+    assert "Probability" in kept[1], kept
+
+
 ################################################################################
 # Reranker Only Tests
 ################################################################################
@@ -171,8 +303,9 @@ def test_search_reranker_only(setup_models, model):
 def test_search(setup_models):
     models = setup_models
     rm = models["intfloat/e5-small-v2"]
+    vs = FaissVS() 
     reranker = models["mixedbread-ai/mxbai-rerank-xsmall-v1"]
-    lotus.settings.configure(rm=rm, reranker=reranker)
+    lotus.settings.configure(rm=rm, vs = vs, reranker=reranker)
 
     data = {
         "Course Name": [

.github/workflows/tests.yml

@@ -155,7 +155,7 @@ jobs:
   rm_test:
     name: Retrieval Model Tests
     runs-on: ubuntu-latest
-    timeout-minutes: 5
+    timeout-minutes: 10
 
     steps:
       - name: Checkout code

examples/op_examples/cluster.py

@@ -2,11 +2,13 @@
 
 import lotus
 from lotus.models import LM, SentenceTransformersRM
+from lotus.vector_store import FaissVS
 
 lm = LM(model="gpt-4o-mini")
 rm = SentenceTransformersRM(model="intfloat/e5-base-v2")
+vs = FaissVS() 
 
-lotus.settings.configure(lm=lm, rm=rm)
+lotus.settings.configure(lm=lm, rm=rm, vs=vs)
 data = {
     "Course Name": [
         "Probability and Random Processes",

examples/op_examples/dedup.py

@@ -2,10 +2,11 @@
 
 import lotus
 from lotus.models import SentenceTransformersRM
+from lotus.vector_store import FaissVS
 
 rm = SentenceTransformersRM(model="intfloat/e5-base-v2")
-
-lotus.settings.configure(rm=rm)
+vs = FaissVS() 
+lotus.settings.configure(rm=rm, vs=vs)
 data = {
     "Text": [
         "Probability and Random Processes",

examples/op_examples/join_cascade.py

@@ -3,11 +3,13 @@
 import lotus
 from lotus.models import LM, SentenceTransformersRM
 from lotus.types import CascadeArgs
+from lotus.vector_store import FaissVS
 
 lm = LM(model="gpt-4o-mini")
 rm = SentenceTransformersRM(model="intfloat/e5-base-v2")
+vs = FaissVS() 
 
-lotus.settings.configure(lm=lm, rm=rm)
+lotus.settings.configure(lm=lm, rm=rm, vs=vs)
 data = {
     "Course Name": [
         "Digital Design and Integrated Circuits",

examples/op_examples/partition.py

@@ -2,11 +2,13 @@
 
 import lotus
 from lotus.models import LM, SentenceTransformersRM
+from lotus.vector_store import FaissVS
 
 lm = LM(max_tokens=2048)
 rm = SentenceTransformersRM(model="intfloat/e5-base-v2")
+vs = FaissVS() 
 
-lotus.settings.configure(lm=lm, rm=rm)
+lotus.settings.configure(lm=lm, rm=rm, vs=vs)
 data = {
     "Course Name": [
         "Probability and Random Processes",

examples/op_examples/search.py

@@ -2,12 +2,14 @@
 
 import lotus
 from lotus.models import LM, CrossEncoderReranker, SentenceTransformersRM
+from lotus.vector_store import FaissVS
 
 lm = LM(model="gpt-4o-mini")
 rm = SentenceTransformersRM(model="intfloat/e5-base-v2")
-reranker = CrossEncoderReranker(model="mixedbread-ai/mxbai-rerank-large-v1")
+reranker = CrossEncoderReranker(model="mixeddbread-ai/mxbai-rerank-large-v1")
+vs = FaissVS() 
 
-lotus.settings.configure(lm=lm, rm=rm, reranker=reranker)
+lotus.settings.configure(lm=lm, rm=rm, reranker=reranker, vs=vs)
 data = {
     "Course Name": [
         "Probability and Random Processes",

examples/op_examples/sim_join.py

@@ -2,11 +2,13 @@
 
 import lotus
 from lotus.models import LM, LiteLLMRM
+from lotus.vector_store import FaissVS
 
 lm = LM(model="gpt-4o-mini")
 rm = LiteLLMRM(model="text-embedding-3-small")
+vs = FaissVS() 
 
-lotus.settings.configure(lm=lm, rm=rm)
+lotus.settings.configure(lm=lm, rm=rm, vs=vs)
 data = {
     "Course Name": [
         "History of the Atlantic World",

lotus/models/colbertv2_rm.py

@@ -6,7 +6,6 @@
 from numpy.typing import NDArray
 from PIL import Image
 
-from lotus.models.rm import RM
 from lotus.types import RMOutput
 
 try:
@@ -16,7 +15,7 @@
     pass
 
 
-class ColBERTv2RM(RM):
+class ColBERTv2RM():
     def __init__(self) -> None:
         self.docs: list[str] | None = None
         self.kwargs: dict[str, Any] = {"doc_maxlen": 300, "nbits": 2}
@@ -46,6 +45,9 @@ def load_index(self, index_dir: str) -> None:
     def get_vectors_from_index(self, index_dir: str, ids: list[int]) -> NDArray[np.float64]:
         raise NotImplementedError("This method is not implemented for ColBERTv2RM")
 
+
+
+   # this should be called in vs.py if it's 
     def __call__(
         self,
         queries: str | Image.Image | list | NDArray[np.float64],