Feature: improve relationship builders for better async and reduced memory utilization (#2077)

- CosineSimilarityBuilder now uses a sharded/chunked similarity
calculation to significantly reduce memory requirements
- CosineSimilarityBuilder and JaccardSimilarityBuilder now leverage
generate_execution_plan to support async iteration over tasks (for
potential future multithreading or improved concurrency)
- Added unit tests

---------

Co-authored-by: Ani <5357586+anistark@users.noreply.github.com>

Author: ahgraber

pyproject.toml

@@ -49,7 +49,7 @@ all = [
     "r2r"
 ]
 
-# Specific integrations  
+# Specific integrations
 tracing = ["langfuse>=3.2.4", "mlflow>=3.1.4"]
 gdrive = [
     "google-api-python-client>=2.178.0",
@@ -70,7 +70,11 @@ dev-minimal = [
     "build>=1.3.0",
 ]
 
-test = []
+# Test only dependencies
+test = [
+    "scipy",
+]
+
 [project.entry-points."ragas.backends"]
 "local/csv" = "ragas.backends.local_csv:LocalCSVBackend"
 "local/jsonl" = "ragas.backends.local_jsonl:LocalJSONLBackend"
@@ -147,7 +151,7 @@ dev = [
     # Core dev tools (shared with minimal)
     "ruff",
     "pyright>=1.1.403",
-    "pre-commit>=4.3.0", 
+    "pre-commit>=4.3.0",
     "pytest",
     "pytest-xdist[psutil]",
     "pytest-asyncio",

src/ragas/testset/transforms/relationship_builders/cosine.py

@@ -12,24 +12,40 @@ class CosineSimilarityBuilder(RelationshipBuilder):
     property_name: str = "embedding"
     new_property_name: str = "cosine_similarity"
     threshold: float = 0.9
+    block_size: int = 1024
+
+    def _block_cosine_similarity(self, i: np.ndarray, j: np.ndarray):
+        """Calculate cosine similarity matrix between two sets of embeddings."""
+        i_norm = i / np.linalg.norm(i, axis=1, keepdims=True)
+        j_norm = j / np.linalg.norm(j, axis=1, keepdims=True)
+        return np.dot(i_norm, j_norm.T)
 
     def _find_similar_embedding_pairs(
         self, embeddings: np.ndarray, threshold: float
     ) -> t.List[t.Tuple[int, int, float]]:
-        # Normalize the embeddings
-        normalized = embeddings / np.linalg.norm(embeddings, axis=1)[:, np.newaxis]
+        """Sharded computation of cosine similarity to find similar pairs."""
 
-        # Calculate cosine similarity matrix
-        similarity_matrix = np.dot(normalized, normalized.T)
-        # Find pairs with similarity >= threshold
-        similar_pairs = np.argwhere(similarity_matrix >= threshold)
+        def process_block(i: int, j: int) -> t.Set[t.Tuple[int, int, float]]:
+            end_i = min(i + self.block_size, n_embeddings)
+            end_j = min(j + self.block_size, n_embeddings)
+            block = self._block_cosine_similarity(
+                embeddings[i:end_i, :], embeddings[j:end_j, :]
+            )
+            similar_idx = np.argwhere(block >= threshold)
+            return {
+                (int(i + ii), int(j + jj), float(block[ii, jj]))
+                for ii, jj in similar_idx
+                if int(i + ii) < int(j + jj)
+            }
 
-        # Filter out self-comparisons and duplicate pairs
-        return [
-            (pair[0], pair[1], similarity_matrix[pair[0], pair[1]])
-            for pair in similar_pairs
-            if pair[0] < pair[1]
-        ]
+        n_embeddings, _dimension = embeddings.shape
+        triplets = set()
+
+        for i in range(0, n_embeddings, self.block_size):
+            for j in range(i, n_embeddings, self.block_size):
+                triplets.update(process_block(i, j))
+
+        return list(triplets)
 
     def _validate_embedding_shapes(self, embeddings: t.List[t.Any]):
         if not embeddings:
@@ -43,40 +59,66 @@ def _validate_embedding_shapes(self, embeddings: t.List[t.Any]):
                 )
 
     async def transform(self, kg: KnowledgeGraph) -> t.List[Relationship]:
-        if self.property_name is None:
-            self.property_name = "embedding"
-
         embeddings = []
         for node in kg.nodes:
             embedding = node.get_property(self.property_name)
             if embedding is None:
                 raise ValueError(f"Node {node.id} has no {self.property_name}")
             embeddings.append(embedding)
-
         self._validate_embedding_shapes(embeddings)
         similar_pairs = self._find_similar_embedding_pairs(
             np.array(embeddings), self.threshold
         )
-
         return [
             Relationship(
                 source=kg.nodes[i],
                 target=kg.nodes[j],
-                type="cosine_similarity",
+                type=self.new_property_name,
                 properties={self.new_property_name: similarity_float},
                 bidirectional=True,
             )
             for i, j, similarity_float in similar_pairs
         ]
 
+    def generate_execution_plan(self, kg: KnowledgeGraph) -> t.List[t.Coroutine]:
+        """
+        Generates a coroutine task for finding similar embedding pairs, which can be scheduled/executed by an Executor.
+        """
+        filtered_kg = self.filter(kg)
+
+        embeddings = []
+        for node in filtered_kg.nodes:
+            embedding = node.get_property(self.property_name)
+            if embedding is None:
+                raise ValueError(f"Node {node.id} has no {self.property_name}")
+            embeddings.append(embedding)
+        self._validate_embedding_shapes(embeddings)
+
+        async def find_and_add_relationships():
+            similar_pairs = self._find_similar_embedding_pairs(
+                np.array(embeddings), self.threshold
+            )
+            for i, j, similarity_float in similar_pairs:
+                rel = Relationship(
+                    source=filtered_kg.nodes[i],
+                    target=filtered_kg.nodes[j],
+                    type=self.new_property_name,
+                    properties={self.new_property_name: similarity_float},
+                    bidirectional=True,
+                )
+                kg.relationships.append(rel)
+
+        return [find_and_add_relationships()]
+
 
 @dataclass
 class SummaryCosineSimilarityBuilder(CosineSimilarityBuilder):
     property_name: str = "summary_embedding"
     new_property_name: str = "summary_cosine_similarity"
     threshold: float = 0.1
+    block_size: int = 1024
 
-    def filter(self, kg: KnowledgeGraph) -> KnowledgeGraph:
+    def _document_summary_filter(self, kg: KnowledgeGraph) -> KnowledgeGraph:
         """
         Filters the knowledge graph to only include nodes with a summary embedding.
         """
@@ -90,22 +132,22 @@ def filter(self, kg: KnowledgeGraph) -> KnowledgeGraph:
         return KnowledgeGraph(nodes=nodes)
 
     async def transform(self, kg: KnowledgeGraph) -> t.List[Relationship]:
+        filtered_kg = self._document_summary_filter(kg)
         embeddings = [
             node.get_property(self.property_name)
-            for node in kg.nodes
+            for node in filtered_kg.nodes
             if node.get_property(self.property_name) is not None
         ]
         if not embeddings:
             raise ValueError(f"No nodes have a valid {self.property_name}")
-        self._validate_embedding_shapes(embeddings)
         similar_pairs = self._find_similar_embedding_pairs(
             np.array(embeddings), self.threshold
         )
         return [
             Relationship(
-                source=kg.nodes[i],
-                target=kg.nodes[j],
-                type="summary_cosine_similarity",
+                source=filtered_kg.nodes[i],
+                target=filtered_kg.nodes[j],
+                type=self.new_property_name,
                 properties={self.new_property_name: similarity_float},
                 bidirectional=True,
             )

src/ragas/testset/transforms/relationship_builders/traditional.py

@@ -1,3 +1,4 @@
+import itertools
 import typing as t
 from collections import Counter
 from dataclasses import dataclass
@@ -19,39 +20,62 @@ def _jaccard_similarity(self, set1: t.Set[str], set2: t.Set[str]) -> float:
         union = len(set1.union(set2))
         return intersection / union if union > 0 else 0.0
 
-    async def transform(self, kg: KnowledgeGraph) -> t.List[Relationship]:
-        if self.property_name is None:
-            self.property_name
-
-        similar_pairs = []
-        for i, node1 in enumerate(kg.nodes):
-            for j, node2 in enumerate(kg.nodes):
-                if i >= j:
-                    continue
-                items1 = node1.get_property(self.property_name)
-                items2 = node2.get_property(self.property_name)
-                if items1 is None or items2 is None:
-                    raise ValueError(
-                        f"Node {node1.id} or {node2.id} has no {self.property_name}"
-                    )
-                if self.key_name is not None:
-                    items1 = items1.get(self.key_name, [])
-                    items2 = items2.get(self.key_name, [])
-                similarity = self._jaccard_similarity(set(items1), set(items2))
-                if similarity >= self.threshold:
-                    similar_pairs.append((i, j, similarity))
+    def _find_similar_embedding_pairs(
+        self, kg: KnowledgeGraph
+    ) -> t.List[t.Tuple[int, int, float]]:
+        """
+        Finds all node index pairs with Jaccard similarity above the threshold.
+        Returns a set of (i, j, similarity) tuples.
+        """
+
+        similar_pairs = set()
+        for (i, node1), (j, node2) in itertools.combinations(enumerate(kg.nodes), 2):
+            items1 = node1.get_property(self.property_name)
+            items2 = node2.get_property(self.property_name)
+            if items1 is None or items2 is None:
+                raise ValueError(
+                    f"Node {node1.id} or {node2.id} has no {self.property_name}"
+                )
+            if self.key_name is not None:
+                items1 = items1.get(self.key_name, [])
+                items2 = items2.get(self.key_name, [])
+            similarity = self._jaccard_similarity(set(items1), set(items2))
+            if similarity >= self.threshold:
+                similar_pairs.add((i, j, similarity))
+        return list(similar_pairs)
 
+    async def transform(self, kg: KnowledgeGraph) -> t.List[Relationship]:
+        similar_pairs = self._find_similar_embedding_pairs(kg)
         return [
             Relationship(
                 source=kg.nodes[i],
                 target=kg.nodes[j],
-                type="jaccard_similarity",
+                type=self.new_property_name,
                 properties={self.new_property_name: similarity_float},
                 bidirectional=True,
             )
             for i, j, similarity_float in similar_pairs
         ]
 
+    def generate_execution_plan(self, kg: KnowledgeGraph) -> t.List[t.Coroutine]:
+        """
+        Generates a coroutine task for finding similar pairs, which can be scheduled/executed by an Executor.
+        """
+
+        async def find_and_add_relationships():
+            similar_pairs = self._find_similar_embedding_pairs(kg)
+            for i, j, similarity_float in similar_pairs:
+                rel = Relationship(
+                    source=kg.nodes[i],
+                    target=kg.nodes[j],
+                    type=self.new_property_name,
+                    properties={self.new_property_name: similarity_float},
+                    bidirectional=True,
+                )
+                kg.relationships.append(rel)
+
+        return [find_and_add_relationships()]
+
 
 @dataclass
 class OverlapScoreBuilder(RelationshipBuilder):
@@ -65,6 +89,7 @@ class OverlapScoreBuilder(RelationshipBuilder):
     def __post_init__(self):
         try:
             from rapidfuzz import distance
+
         except ImportError:
             raise ImportError(
                 "rapidfuzz is required for string distance. Please install it using `pip install rapidfuzz`"