Repository: Marker-Inc-Korea/AutoRAG

Issue:
Add threshold cutoff module

Discussion:

Patch:
diff --git a/autorag/nodes/passagefilter/__init__.py b/autorag/nodes/passagefilter/__init__.py
index 0ecb3c7..dce8ff9 100644
--- a/autorag/nodes/passagefilter/__init__.py
+++ b/autorag/nodes/passagefilter/__init__.py
@@ -1,4 +1,5 @@
-from .percentile_cutoff import similarity_percentile_cutoff
 from .pass_passage_filter import pass_passage_filter
 from .recency import recency_filter
-from .threshold_cutoff import similarity_threshold_cutoff
+from .similarity_percentile_cutoff import similarity_percentile_cutoff
+from .similarity_threshold_cutoff import similarity_threshold_cutoff
+from .threshold_cutoff import threshold_cutoff
diff --git a/autorag/nodes/passagefilter/percentile_cutoff.py b/autorag/nodes/passagefilter/similarity_percentile_cutoff.py
similarity index 97%
rename from autorag/nodes/passagefilter/percentile_cutoff.py
rename to autorag/nodes/passagefilter/similarity_percentile_cutoff.py
index b63cb19..644e7a4 100644
--- a/autorag/nodes/passagefilter/percentile_cutoff.py
+++ b/autorag/nodes/passagefilter/similarity_percentile_cutoff.py
@@ -5,7 +5,7 @@ import torch.cuda
 
 from autorag.evaluate.metric.util import calculate_cosine_similarity
 from autorag.nodes.passagefilter.base import passage_filter_node
-from autorag.nodes.passagefilter.threshold_cutoff import embedding_query_content
+from autorag.nodes.passagefilter.similarity_threshold_cutoff import embedding_query_content
 
 
 @passage_filter_node
diff --git a/autorag/nodes/passagefilter/similarity_threshold_cutoff.py b/autorag/nodes/passagefilter/similarity_threshold_cutoff.py
new file mode 100644
index 0000000..b5461d4
--- /dev/null
+++ b/autorag/nodes/passagefilter/similarity_threshold_cutoff.py
@@ -0,0 +1,88 @@
+import itertools
+from typing import List, Tuple, Optional
+
+import numpy as np
+import torch.cuda
+
+from autorag import embedding_models
+from autorag.evaluate.metric.util import calculate_cosine_similarity
+from autorag.nodes.passagefilter.base import passage_filter_node
+from autorag.utils.util import reconstruct_list
+
+
+@passage_filter_node
+def similarity_threshold_cutoff(queries: List[str], contents_list: List[List[str]],
+                                scores_list: List[List[float]], ids_list: List[List[str]],
+                                threshold: float, embedding_model: Optional[str] = None,
+                                batch: int = 128,
+                                ) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+    """
+    Re-calculate each content's similarity with the query and filter out the contents that are below the threshold.
+    If all contents are filtered, keep the only one highest similarity content.
+    This is a filter and does not override scores.
+    The output of scores is not coming from query-content similarity.
+
+    :param queries: The list of queries to use for filtering
+    :param contents_list: The list of lists of contents to filter
+    :param scores_list: The list of lists of scores retrieved
+    :param ids_list: The list of lists of ids retrieved
+    :param threshold: The threshold to cut off
+    :param embedding_model: The embedding model to use for calculating similarity
+        Default is OpenAIEmbedding.
+    :param batch: The number of queries to be processed in a batch
+        Default is 128.
+    :return: Tuple of lists containing the filtered contents, ids, and scores
+    """
+    query_embeddings, content_embeddings = embedding_query_content(queries, contents_list, embedding_model, batch)
+
+    remain_indices = list(map(lambda x: similarity_threshold_cutoff_pure(x[0], x[1], threshold),
+                              zip(query_embeddings, content_embeddings)))
+
+    remain_content_list = list(map(lambda c, idx: [c[i] for i in idx], contents_list, remain_indices))
+    remain_scores_list = list(map(lambda s, idx: [s[i] for i in idx], scores_list, remain_indices))
+    remain_ids_list = list(map(lambda _id, idx: [_id[i] for i in idx], ids_list, remain_indices))
+
+    del embedding_model
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    return remain_content_list, remain_ids_list, remain_scores_list
+
+
+def similarity_threshold_cutoff_pure(query_embedding: str,
+                                     content_embeddings: List[List[float]],
+                                     threshold: float) -> List[int]:
+    """
+    Return indices that have to remain.
+    Return at least one index if there is nothing to remain.
+
+    :param query_embedding: Query embedding
+    :param content_embeddings: Each content embedding
+    :param threshold: The threshold to cut off
+    :return: Indices to remain at the contents
+    """
+
+    similarities = np.array(list(map(lambda x: calculate_cosine_similarity(query_embedding, x),
+                                     content_embeddings)))
+    result = np.where(similarities >= threshold)[0].tolist()
+    if len(result) > 0:
+        return result
+    return [np.argmax(similarities)]
+
+
+def embedding_query_content(queries: List[str], contents_list: List[List[str]],
+                            embedding_model: Optional[str] = None, batch: int = 128):
+    if embedding_model is None:
+        embedding_model = embedding_models['openai']
+    else:
+        embedding_model = embedding_models[embedding_model]
+
+    # Embedding using batch
+    embedding_model.embed_batch_size = batch
+    query_embeddings = embedding_model.get_text_embedding_batch(queries)
+
+    content_lengths = list(map(len, contents_list))
+    content_embeddings_flatten = embedding_model.get_text_embedding_batch(list(
+        itertools.chain.from_iterable(contents_list)))
+    content_embeddings = reconstruct_list(content_embeddings_flatten, content_lengths)
+    return query_embeddings, content_embeddings
diff --git a/autorag/nodes/passagefilter/threshold_cutoff.py b/autorag/nodes/passagefilter/threshold_cutoff.py
index b5461d4..1369344 100644
--- a/autorag/nodes/passagefilter/threshold_cutoff.py
+++ b/autorag/nodes/passagefilter/threshold_cutoff.py
@@ -1,88 +1,53 @@
-import itertools
-from typing import List, Tuple, Optional
+from typing import List, Tuple
 
-import numpy as np
-import torch.cuda
-
-from autorag import embedding_models
-from autorag.evaluate.metric.util import calculate_cosine_similarity
 from autorag.nodes.passagefilter.base import passage_filter_node
-from autorag.utils.util import reconstruct_list
 
 
 @passage_filter_node
-def similarity_threshold_cutoff(queries: List[str], contents_list: List[List[str]],
-                                scores_list: List[List[float]], ids_list: List[List[str]],
-                                threshold: float, embedding_model: Optional[str] = None,
-                                batch: int = 128,
-                                ) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+def threshold_cutoff(queries: List[str], contents_list: List[List[str]],
+                     scores_list: List[List[float]], ids_list: List[List[str]],
+                     threshold: float, reverse: bool = False,
+                     ) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
     """
-    Re-calculate each content's similarity with the query and filter out the contents that are below the threshold.
-    If all contents are filtered, keep the only one highest similarity content.
-    This is a filter and does not override scores.
-    The output of scores is not coming from query-content similarity.
-
-    :param queries: The list of queries to use for filtering
-    :param contents_list: The list of lists of contents to filter
-    :param scores_list: The list of lists of scores retrieved
-    :param ids_list: The list of lists of ids retrieved
-    :param threshold: The threshold to cut off
-    :param embedding_model: The embedding model to use for calculating similarity
-        Default is OpenAIEmbedding.
-    :param batch: The number of queries to be processed in a batch
-        Default is 128.
-    :return: Tuple of lists containing the filtered contents, ids, and scores
+    Filters the contents, scores, and ids based on a previous result's score.
+    Keeps at least one item per query if all scores are below the threshold.
+
+    :param queries: List of query strings (not used in the current implementation).
+    :param contents_list: List of content strings for each query.
+    :param scores_list: List of scores for each content.
+    :param ids_list: List of ids for each content.
+    :param threshold: The minimum score to keep an item.
+    :param reverse: If True, the lower the score, the better.
+        Default is False.
+    :return: Filtered lists of contents, ids, and scores.
     """
-    query_embeddings, content_embeddings = embedding_query_content(queries, contents_list, embedding_model, batch)
-
-    remain_indices = list(map(lambda x: similarity_threshold_cutoff_pure(x[0], x[1], threshold),
-                              zip(query_embeddings, content_embeddings)))
+    remain_indices = list(map(lambda x: threshold_cutoff_pure(x, threshold, reverse), scores_list))
 
     remain_content_list = list(map(lambda c, idx: [c[i] for i in idx], contents_list, remain_indices))
     remain_scores_list = list(map(lambda s, idx: [s[i] for i in idx], scores_list, remain_indices))
     remain_ids_list = list(map(lambda _id, idx: [_id[i] for i in idx], ids_list, remain_indices))
 
-    del embedding_model
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
     return remain_content_list, remain_ids_list, remain_scores_list
 
 
-def similarity_threshold_cutoff_pure(query_embedding: str,
-                                     content_embeddings: List[List[float]],
-                                     threshold: float) -> List[int]:
+def threshold_cutoff_pure(scores_list: List[float],
+                          threshold: float,
+                          reverse: bool = False) -> List[int]:
     """
     Return indices that have to remain.
     Return at least one index if there is nothing to remain.
 
-    :param query_embedding: Query embedding
-    :param content_embeddings: Each content embedding
+    :param scores_list: Each score
     :param threshold: The threshold to cut off
+    :param reverse: If True, the lower the score, the better
+        Default is False.
     :return: Indices to remain at the contents
     """
-
-    similarities = np.array(list(map(lambda x: calculate_cosine_similarity(query_embedding, x),
-                                     content_embeddings)))
-    result = np.where(similarities >= threshold)[0].tolist()
-    if len(result) > 0:
-        return result
-    return [np.argmax(similarities)]
-
-
-def embedding_query_content(queries: List[str], contents_list: List[List[str]],
-                            embedding_model: Optional[str] = None, batch: int = 128):
-    if embedding_model is None:
-        embedding_model = embedding_models['openai']
+    if reverse:
+        remain_indices = [i for i, score in enumerate(scores_list) if score <= threshold]
+        default_index = scores_list.index(min(scores_list))
     else:
-        embedding_model = embedding_models[embedding_model]
-
-    # Embedding using batch
-    embedding_model.embed_batch_size = batch
-    query_embeddings = embedding_model.get_text_embedding_batch(queries)
+        remain_indices = [i for i, score in enumerate(scores_list) if score >= threshold]
+        default_index = scores_list.index(max(scores_list))
 
-    content_lengths = list(map(len, contents_list))
-    content_embeddings_flatten = embedding_model.get_text_embedding_batch(list(
-        itertools.chain.from_iterable(contents_list)))
-    content_embeddings = reconstruct_list(content_embeddings_flatten, content_lengths)
-    return query_embeddings, content_embeddings
+    return remain_indices if remain_indices else [default_index]
diff --git a/autorag/support.py b/autorag/support.py
index a832491..e552438 100644
--- a/autorag/support.py
+++ b/autorag/support.py
@@ -47,6 +47,7 @@ def get_support_modules(module_name: str) -> Callable:
         'similarity_threshold_cutoff': ('autorag.nodes.passagefilter', 'similarity_threshold_cutoff'),
         'similarity_percentile_cutoff': ('autorag.nodes.passagefilter', 'similarity_percentile_cutoff'),
         'recency_filter': ('autorag.nodes.passagefilter', 'recency_filter'),
+        'threshold_cutoff': ('autorag.nodes.passagefilter', 'threshold_cutoff'),
         # passage_compressor
         'tree_summarize': ('autorag.nodes.passagecompressor', 'tree_summarize'),
         'pass_compressor': ('autorag.nodes.passagecompressor', 'pass_compressor'),
diff --git a/docs/source/api_spec/autorag.nodes.passagefilter.rst b/docs/source/api_spec/autorag.nodes.passagefilter.rst
index 9db5c7a..a47da9e 100644
--- a/docs/source/api_spec/autorag.nodes.passagefilter.rst
+++ b/docs/source/api_spec/autorag.nodes.passagefilter.rst
@@ -20,14 +20,6 @@ autorag.nodes.passagefilter.pass\_passage\_filter module
    :undoc-members:
    :show-inheritance:
 
-autorag.nodes.passagefilter.percentile\_cutoff module
------------------------------------------------------
-
-.. automodule:: autorag.nodes.passagefilter.percentile_cutoff
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
 autorag.nodes.passagefilter.recency module
 ------------------------------------------
 
@@ -44,6 +36,22 @@ autorag.nodes.passagefilter.run module
    :undoc-members:
    :show-inheritance:
 
+autorag.nodes.passagefilter.similarity\_percentile\_cutoff module
+-----------------------------------------------------------------
+
+.. automodule:: autorag.nodes.passagefilter.similarity_percentile_cutoff
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+autorag.nodes.passagefilter.similarity\_threshold\_cutoff module
+----------------------------------------------------------------
+
+.. automodule:: autorag.nodes.passagefilter.similarity_threshold_cutoff
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 autorag.nodes.passagefilter.threshold\_cutoff module
 ----------------------------------------------------
 
diff --git a/docs/source/nodes/passage_filter/passage_filter.md b/docs/source/nodes/passage_filter/passage_filter.md
index 15071e1..d373ee0 100644
--- a/docs/source/nodes/passage_filter/passage_filter.md
+++ b/docs/source/nodes/passage_filter/passage_filter.md
@@ -52,4 +52,5 @@ maxdepth: 1
 similarity_threshold_cutoff.md
 similarity_percentile_cutoff.md
 recency_filter.md
+threshold_cutoff.md
 ```
diff --git a/docs/source/nodes/passage_filter/similarity_threshold_cutoff.md b/docs/source/nodes/passage_filter/similarity_threshold_cutoff.md
index eaa380b..b137489 100644
--- a/docs/source/nodes/passage_filter/similarity_threshold_cutoff.md
+++ b/docs/source/nodes/passage_filter/similarity_threshold_cutoff.md
@@ -4,6 +4,8 @@ This module is inspired by
 LlamaIndex ['Sentence Embedding Optimizer'](https://docs.llamaindex.ai/en/stable/examples/node_postprocessor/OptimizerDemo/).
 Re-calculate each content's similarity with the query and filter out the contents that are below the threshold.
 
+📣 Keeps at least one item per query if all scores are below the threshold.
+
 ## **Module Parameters**
 
 - **threshold** : The threshold value to filter out the contents.
diff --git a/docs/source/nodes/passage_filter/threshold_cutoff.md b/docs/source/nodes/passage_filter/threshold_cutoff.md
new file mode 100644
index 0000000..40e77ba
--- /dev/null
+++ b/docs/source/nodes/passage_filter/threshold_cutoff.md
@@ -0,0 +1,24 @@
+# Threshold Cutoff
+
+This module is inspired by
+our [similarity threshold cutoff](https://marker-inc-korea.github.io/AutoRAG/nodes/passage_filter/similarity_threshold_cutoff.html)
+module.
+Filters the contents, scores, and ids based on a **previous result's scores**.
+
+📣 Keeps at least one item per query if all scores are below the threshold.
+
+## **Module Parameters**
+
+- **threshold** : The threshold value to filter out the contents.
+  If the score is below the threshold, the content will be filtered out.
+  This is essential to run the module, so you have to set this parameter.
+- **reverse** : If True, the lower the score, the better.
+  Default is False.
+
+## **Example config.yaml**
+
+```yaml
+modules:
+  - module_type: threshold_cutoff
+    threshold: 0.85
+```
diff --git a/sample_config/full.yaml b/sample_config/full.yaml
index b15522d..cbb7df3 100644
--- a/sample_config/full.yaml
+++ b/sample_config/full.yaml
@@ -88,8 +88,12 @@ node_lines:
         - module_type: pass_passage_filter
         - module_type: similarity_threshold_cutoff
           threshold: 0.85
+        - module_type: similarity_percentile_cutoff
+          percentile: 0.6
         - module_type: recency_filter
           threshold: 2015-01-01
+        - module_type: threshold_cutoff
+          threshold: 0.85
     - node_type: passage_compressor
       strategy:
         metrics: [retrieval_token_f1, retrieval_token_recall, retrieval_token_precision]


Test patch:
diff --git a/tests/autorag/nodes/passagefilter/test_percentile_cutoff.py b/tests/autorag/nodes/passagefilter/test_similarity_percentile_cutoff.py
similarity index 100%
rename from tests/autorag/nodes/passagefilter/test_percentile_cutoff.py
rename to tests/autorag/nodes/passagefilter/test_similarity_percentile_cutoff.py
diff --git a/tests/autorag/nodes/passagefilter/test_similarity_threshold_cutoff.py b/tests/autorag/nodes/passagefilter/test_similarity_threshold_cutoff.py
new file mode 100644
index 0000000..9a42782
--- /dev/null
+++ b/tests/autorag/nodes/passagefilter/test_similarity_threshold_cutoff.py
@@ -0,0 +1,17 @@
+from autorag.nodes.passagefilter import similarity_threshold_cutoff
+from tests.autorag.nodes.passagefilter.test_passage_filter_base import queries_example, contents_example, \
+    scores_example, ids_example, base_passage_filter_test, project_dir, previous_result, base_passage_filter_node_test
+
+
+def test_similarity_threshold_cutoff():
+    original_cutoff = similarity_threshold_cutoff.__wrapped__
+    contents, ids, scores = original_cutoff(
+        queries_example, contents_example, scores_example, ids_example, threshold=0.85,
+        embedding_model='openai_embed_3_large', batch=64)
+    base_passage_filter_test(contents, ids, scores)
+
+
+def test_similarity_threshold_cutoff_node():
+    result_df = similarity_threshold_cutoff(
+        project_dir=project_dir, previous_result=previous_result, threshold=0.9)
+    base_passage_filter_node_test(result_df)
diff --git a/tests/autorag/nodes/passagefilter/test_threshold_cutoff.py b/tests/autorag/nodes/passagefilter/test_threshold_cutoff.py
index 9a42782..59d1d91 100644
--- a/tests/autorag/nodes/passagefilter/test_threshold_cutoff.py
+++ b/tests/autorag/nodes/passagefilter/test_threshold_cutoff.py
@@ -1,17 +1,25 @@
-from autorag.nodes.passagefilter import similarity_threshold_cutoff
+from autorag.nodes.passagefilter import threshold_cutoff
 from tests.autorag.nodes.passagefilter.test_passage_filter_base import queries_example, contents_example, \
     scores_example, ids_example, base_passage_filter_test, project_dir, previous_result, base_passage_filter_node_test
 
 
-def test_similarity_threshold_cutoff():
-    original_cutoff = similarity_threshold_cutoff.__wrapped__
+def test_threshold_cutoff():
+    original_cutoff = threshold_cutoff.__wrapped__
     contents, ids, scores = original_cutoff(
-        queries_example, contents_example, scores_example, ids_example, threshold=0.85,
-        embedding_model='openai_embed_3_large', batch=64)
+        queries_example, contents_example, scores_example, ids_example, threshold=0.6)
     base_passage_filter_test(contents, ids, scores)
 
 
-def test_similarity_threshold_cutoff_node():
-    result_df = similarity_threshold_cutoff(
+def test_threshold_cutoff_reverse():
+    original_cutoff = threshold_cutoff.__wrapped__
+    contents, ids, scores = original_cutoff(
+        queries_example, contents_example, scores_example, ids_example, threshold=0.4, reverse=True)
+    base_passage_filter_test(contents, ids, scores)
+    assert scores[0] == [0.1, 0.1]
+    assert contents[0] == ["NomaDamas is Great Team", "havertz is suck at soccer"]
+
+
+def test_threshold_cutoff_node():
+    result_df = threshold_cutoff(
         project_dir=project_dir, previous_result=previous_result, threshold=0.9)
     base_passage_filter_node_test(result_df)