Recommender endpoint

requirements.dev.txt

@@ -1,4 +1,4 @@
-# tensor search, idential to requirements.txt:
+# tensor search, identical to requirements.txt:
 requests==2.28.1
 anyio==3.7.1
 fastapi==0.86.0

requirements.txt

@@ -12,4 +12,5 @@ urllib3==1.26.0
 pydantic==1.10.11
 httpx==0.25.0
 semver==3.0.2
-memory-profiler==0.61.0
+scipy==1.10.1
+memory-profiler==0.61.0

src/marqo/case_insensitive_enum.py

@@ -0,0 +1,12 @@
+from enum import Enum
+from typing import Optional
+
+
+class CaseInsensitiveEnum(Enum):
+    @classmethod
+    def _missing_(cls, value: str) -> Optional['CaseInsensitiveEnum']:
+        value = value.lower()
+        for member in cls:
+            if member.value.lower() == value:
+                return member
+        return None

src/marqo/config.py

@@ -1,17 +1,17 @@
 from typing import Optional, Union
 
+from marqo.core.document.document import Document
 from marqo.core.index_management.index_management import IndexManagement
 from marqo.core.monitoring.monitoring import Monitoring
+from marqo.core.search.recommender import Recommender
 from marqo.tensor_search import enums
-from marqo.core.document.document import Document
 from marqo.vespa.vespa_client import VespaClient
 
 
 class Config:
     def __init__(
             self,
             vespa_client: VespaClient,
-            index_management: IndexManagement,
             timeout: Optional[int] = None,
             backend: Optional[Union[enums.SearchDb, str]] = None,
     ) -> None:
@@ -27,9 +27,10 @@ def __init__(
         self.backend = backend if backend is not None else enums.SearchDb.vespa
 
         # Initialize Core layer dependencies
-        self.index_management = index_management
-        self.monitoring = Monitoring(vespa_client, index_management)
-        self.document = Document(vespa_client, index_management)
+        self.index_management = IndexManagement(vespa_client)
+        self.monitoring = Monitoring(vespa_client, self.index_management)
+        self.document = Document(vespa_client, self.index_management)
+        self.recommender = Recommender(vespa_client, self.index_management)
 
     def set_is_remote(self, vespa_client: VespaClient):
         local_host_markers = ["localhost", "0.0.0.0", "127.0.0.1"]

src/marqo/core/models/interpolation_method.py

@@ -0,0 +1,7 @@
+from marqo.case_insensitive_enum import CaseInsensitiveEnum
+
+
+class InterpolationMethod(str, CaseInsensitiveEnum):
+    LERP = "lerp"
+    NLERP = "nlerp"
+    SLERP = "slerp"

src/marqo/core/search/recommender.py

@@ -0,0 +1,179 @@
+from timeit import default_timer as timer
+from typing import Dict, List, Union, Optional
+
+from marqo.core.exceptions import InvalidFieldNameError
+from marqo.core.index_management.index_management import IndexManagement
+from marqo.core.models import MarqoIndex
+from marqo.core.models.interpolation_method import InterpolationMethod
+from marqo.core.models.marqo_index import IndexType
+from marqo.core.utils.vector_interpolation import from_interpolation_method
+from marqo.exceptions import InvalidArgumentError
+from marqo.tensor_search.models.score_modifiers_object import ScoreModifier
+from marqo.tensor_search.models.search import SearchContext, SearchContextTensor
+from marqo.vespa.vespa_client import VespaClient
+
+
+class Recommender:
+    def __init__(self, vespa_client: VespaClient, index_management: IndexManagement):
+        self.vespa_client = vespa_client
+        self.index_management = index_management
+
+    def recommend(self,
+                  index_name: str,
+                  documents: Union[List[str], Dict[str, float]],
+                  tensor_fields: Optional[List[str]] = None,
+                  interpolation_method: Optional[InterpolationMethod] = None,
+                  exclude_input_documents: bool = True,
+                  result_count: int = 3,
+                  offset: int = 0,
+                  highlights: bool = True,
+                  ef_search: Optional[int] = None,
+                  approximate: Optional[bool] = None,
+                  searchable_attributes: Optional[List[str]] = None,
+                  verbose: int = 0,
+                  reranker: Union[str, Dict] = None,
+                  filter: str = None,
+                  attributes_to_retrieve: Optional[List[str]] = None,
+                  score_modifiers: Optional[ScoreModifier] = None
+                  ):
+        """
+        Recommend documents similar to the provided documents.
+
+        Args:
+            index_name: Name of the index to search
+            documents: A list of document IDs or a dictionary where the keys are document IDs and the values are weights
+            tensor_fields: List of tensor fields to use for recommendation
+            interpolation_method: Interpolation method to use for combining vectors
+            exclude_input_documents: Whether to exclude the input documents from the search results
+            result_count: Number of results to return
+            offset: Offset of the first result
+            highlights: Whether to include highlights in the results
+            ef_search: ef_search parameter for HNSW search
+            approximate: Whether to use approximate search
+            searchable_attributes: List of attributes to search in
+            verbose: Verbosity level
+            reranker: Reranker to use
+            filter: Filter string
+            attributes_to_retrieve: List of attributes to retrieve
+            score_modifiers: Score modifiers to apply
+        """
+        # TODO - Extract search and get_docs from tensor_search and refactor this
+        # TODO - The dependence on Config in tensor_search is bad design. Refactor to require specific dependencies
+        from marqo import config
+        from marqo.tensor_search import tensor_search
+        from marqo.tensor_search import index_meta_cache
+
+        if documents is None or len(documents) == 0:
+            raise InvalidArgumentError('No document IDs provided')
+
+        marqo_index = index_meta_cache.get_index(config.Config(self.vespa_client), index_name=index_name)
+
+        if interpolation_method is None:
+            interpolation_method = self._get_default_interpolation_method(marqo_index)
+
+        vector_interpolation = from_interpolation_method(interpolation_method)
+
+        if marqo_index.type == IndexType.Structured:
+            # Validate tensor field names
+            if tensor_fields is not None:
+                valid_tensor_fields = marqo_index.tensor_field_map.keys()
+                for tensor_field in tensor_fields:
+                    if tensor_field not in valid_tensor_fields:
+                        raise InvalidFieldNameError(f'Tensor field "{tensor_field}" not found in index "{index_name}". '
+                                                    f'Available tensor fields: {", ".join(valid_tensor_fields)}')
+
+        if isinstance(documents, dict):
+            document_ids = list(documents.keys())
+        else:
+            document_ids = documents
+
+        t0 = timer()
+
+        marqo_documents = tensor_search.get_documents_by_ids(
+            config.Config(self.vespa_client),
+            index_name, document_ids, show_vectors=True
+        )
+
+        # Make sure all documents were found
+        not_found = []
+        for document in marqo_documents['results']:
+            if not document['_found']:
+                not_found.append(document['_id'])
+
+        if len(not_found) > 0:
+            raise InvalidArgumentError(f'The following document IDs were not found: {", ".join(not_found)}')
+
+        doc_vectors: Dict[str, List[List[float]]] = {}
+        docs_without_vectors = []
+        for document in marqo_documents['results']:
+            vectors: List[List[float]] = []
+            for tensor_facet in document['_tensor_facets']:
+                field = list(tensor_facet.keys())[0]
+                if tensor_fields is None or field in tensor_fields:
+                    vectors.append(tensor_facet['_embedding'])
+
+            doc_vectors[document['_id']] = vectors
+
+            if len(vectors) == 0:
+                docs_without_vectors.append(document['_id'])
+
+        if len(docs_without_vectors) > 0:
+            raise InvalidArgumentError(
+                f'The following documents do not have embeddings: {", ".join(docs_without_vectors)}'
+            )
+
+        vectors: List[List[float]] = []
+        weights: List[float] = []
+
+        for document_id, vector_list in doc_vectors.items():
+            if isinstance(documents, dict):
+                weight = documents[document_id]
+            else:
+                weight = 1
+            vectors.extend(vector_list)
+            weights.extend([weight] * len(vector_list))
+
+        interpolated_vector = vector_interpolation.interpolate(
+            vectors, weights
+        )
+
+        if exclude_input_documents:
+            recommend_filter = self._get_exclusion_filter(document_ids, filter)
+        else:
+            recommend_filter = filter
+
+        results = tensor_search.search(
+            config.Config(self.vespa_client),
+            index_name,
+            text=None,
+            context=SearchContext(tensor=[SearchContextTensor(vector=interpolated_vector, weight=1)]),
+            result_count=result_count,
+            offset=offset,
+            highlights=highlights,
+            ef_search=ef_search,
+            approximate=approximate,
+            searchable_attributes=searchable_attributes,
+            verbose=verbose,
+            reranker=reranker,
+            filter=recommend_filter,
+            attributes_to_retrieve=attributes_to_retrieve,
+            device=device,
+            score_modifiers=score_modifiers,
+            processing_start=t0
+        )
+
+        return results
+
+    def _get_default_interpolation_method(self, marqo_index: MarqoIndex) -> InterpolationMethod:
+        if marqo_index.normalize_embeddings:
+            return InterpolationMethod.SLERP
+        else:
+            return InterpolationMethod.LERP
+
+    def _get_exclusion_filter(self, documents: List[str], user_filter: Optional[str]) -> str:
+        not_in = 'NOT (' + ' OR '.join([f'_id:({doc})' for doc in documents]) + ')'
+
+        if user_filter is not None and user_filter.strip() != '':
+            return f'({user_filter}) AND {not_in}'
+        else:
+            return not_in

src/marqo/core/utils/vector_interpolation.py

@@ -0,0 +1,126 @@
+import abc
+import math
+from enum import Enum
+from typing import List
+
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+from scipy.spatial.transform import Slerp
+
+from marqo.core.models.interpolation_method import InterpolationMethod
+from marqo.exceptions import InternalError
+
+
+class VectorInterpolation(abc.ABC):
+    @abc.abstractmethod
+    def interpolate(self, vectors: List[List[float]], weights: List[float] = None) -> List[float]:
+        """
+        Interpolates a list of vectors using the given weights.
+
+        Args:
+            vectors: A list of vectors to interpolate.
+            weights: A list of weights to use for the interpolation. If None, the interpolation will be done using
+            euqual weights for all vectors.
+
+        Returns:
+            The interpolated vector.
+        """
+        pass
+
+
+def from_interpolation_method(method: InterpolationMethod):
+    if method == InterpolationMethod.SLERP:
+        return Slerp()
+    elif method == InterpolationMethod.NLERP:
+        return Nlerp()
+    elif method == InterpolationMethod.LERP:
+        return Lerp()
+    else:
+        raise InternalError(f'Unknown interpolation method: {method}')
+
+
+class Lerp(VectorInterpolation):
+    def interpolate(self, vectors: List[List[float]], weights: List[float] = None) -> List[float]:
+        if len(vectors) < 1:
+            raise ValueError('Cannot interpolate an empty list of vectors')
+
+        if not weights:
+            weights = [1] * len(vectors)
+        weight_sum = sum(weights)
+        result = [0] * len(vectors[0])
+        for vector, weight in zip(vectors, weights):
+            for i, value in enumerate(vector):
+                result[i] += (weight / weight_sum) * value
+        return result
+
+
+class Nlerp(VectorInterpolation):
+    def interpolate(self, vectors: List[List[float]], weights: List[float] = None) -> List[float]:
+        lerp_result = Lerp().interpolate(vectors, weights)
+        length = math.sqrt(sum(x ** 2 for x in lerp_result))
+        return [x / length for x in lerp_result]
+
+
+class Slerp(VectorInterpolation):
+    class Method(Enum):
+        Sequential = 0
+        Hierarchical = 1
+
+    def __init__(self, method: Method = Method.Hierarchical):
+        self.method = method
+        self._slerp_impl = self._slerp
+
+    def _slerp_scipy(self, v1: List[float], v2: List[float], t: float) -> List[float]:
+        # Convert vectors to quaternions (assuming vectors are 3D)
+        q1 = np.array([0] + v1)
+        q2 = np.array([0] + v2)
+
+        # Create a Slerp object and interpolate
+        slerp = Slerp([0, 1], R.from_quat([q1, q2]))
+        interpolated_quaternion = slerp([t]).as_quat()[0]
+
+        # Return the vector part of the interpolated quaternion
+        return interpolated_quaternion[1:].tolist()
+
+    def _slerp(self, v0: List[float], v1: List[float], t: float) -> List[float]:
+        dot = sum(a * b for a, b in zip(v0, v1))
+        theta = math.acos(dot) * t
+        relative_vec = [b - a * dot for a, b in zip(v0, v1)]
+        return [math.cos(theta) * a + math.sin(theta) * b for a, b in zip(v0, relative_vec)]
+
+    def _interpolate_sequential(self, vectors: List[List[float]], weights: List[float]) -> List[float]:
+        result = vectors[0]
+        for i in range(1, len(vectors)):
+            result = self._slerp_impl(result, vectors[i], weights[i])
+        return result
+
+    def _interpolate_hierarchical(self, vectors: List[List[float]], weights: List[float]) -> List[float]:
+        while len(vectors) > 1:
+            result = []
+            new_weights = []
+            for i in range(0, len(vectors), 2):
+                if i + 1 == len(vectors):
+                    # if there is an odd number of vectors, just append the last one
+                    result.append(vectors[i])
+                    new_weights.append(weights[i])
+                    continue
+
+                result.append(
+                    self._slerp_impl(vectors[i], vectors[i + 1], weights[i + 1] / (weights[i] + weights[i + 1]))
+                )
+                new_weights.append((weights[i] + weights[i + 1]) / 2)
+            vectors = result
+            weights = new_weights
+
+        return vectors[0]
+
+    def interpolate(self, vectors: List[List[float]], weights: List[float] = None) -> List[float]:
+        if len(vectors) < 1:
+            raise ValueError('Cannot interpolate an empty list of vectors')
+
+        if self.method == self.Method.Sequential:
+            return self._interpolate_sequential(vectors, weights)
+        elif self.method == self.Method.Hierarchical:
+            return self._interpolate_hierarchical(vectors, weights)
+        else:
+            raise InternalError(f'Unknown interpolation method: {self.method}')