From a4eb64d88c8c7985b6985b733254ff2f4bb21dd0 Mon Sep 17 00:00:00 2001
From: Jakub Zakrzewski <jzakrzewski@nvidia.com>
Date: Fri, 3 Oct 2025 13:09:33 +0200
Subject: [PATCH] feat: add Cohere API support to rankings converter

Extend the RankingsConverter to support Cohere's reranking API format,
which uses flat "documents" arrays instead of nested passage objects.

This is enabled by setting rankings="cohere" in the configuration,
similar to the existing TEI support (rankings="tei").
---
 .../inputs/converters/rankings_converter.py   | 18 +++++++-
 .../test_rankings_converter.py                | 46 +++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/genai-perf/genai_perf/inputs/converters/rankings_converter.py b/genai-perf/genai_perf/inputs/converters/rankings_converter.py
index cf93233a..3af74ce3 100644
--- a/genai-perf/genai_perf/inputs/converters/rankings_converter.py
+++ b/genai-perf/genai_perf/inputs/converters/rankings_converter.py
@@ -71,6 +71,13 @@ def convert(
             if self._is_rankings_tei():
                 passages = passage_entry.texts
                 payload = {"query": query, "texts": passages}
+            elif self._is_rankings_cohere():
+                documents = passage_entry.texts
+                payload = {
+                    "query": query,
+                    "documents": documents,
+                    "model": model_name,
+                }
             else:
                 passages = [{"text": p} for p in passage_entry.texts if p is not None]
                 payload = {
@@ -92,8 +99,17 @@ def _is_rankings_tei(self) -> bool:
             return True
         return False
 
+    def _is_rankings_cohere(self) -> bool:
+        """
+        Check if user specified that they are using the Cohere API
+        for ranking models
+        """
+        if self.config.input.extra and self.config.input.extra.get("rankings") == "cohere":
+            return True
+        return False
+
     def _add_request_params(self, payload: Dict, optional_data: Dict[Any, Any]) -> None:
         if self.config.input.extra:
             for key, value in self.config.input.extra.items():
-                if not (key == "rankings" and value == "tei"):
+                if not (key == "rankings" and value in ["tei", "cohere"]):
                     payload[key] = value
diff --git a/genai-perf/tests/test_converters/test_rankings_converter.py b/genai-perf/tests/test_converters/test_rankings_converter.py
index 80250483..a0efec6f 100644
--- a/genai-perf/tests/test_converters/test_rankings_converter.py
+++ b/genai-perf/tests/test_converters/test_rankings_converter.py
@@ -239,6 +239,52 @@ def test_convert_huggingface_tei(self):
 
         assert result == expected_result
 
+    def test_convert_cohere(self):
+        generic_dataset = self.create_generic_dataset(
+            queries_data=[["query 1"], ["query 2"]],
+            passages_data=[["passage 1", "passage 2"], ["passage 3", "passage 4"]],
+        )
+
+        extra_inputs = {
+            "rankings": "cohere",
+            "additional_key": "additional_value",
+        }
+
+        config = ConfigCommand({"model_name": "test_model"})
+        config.endpoint.model_selection_strategy = ModelSelectionStrategy.ROUND_ROBIN
+        config.endpoint.output_format = OutputFormat.RANKINGS
+        config.input.extra = extra_inputs
+
+        rankings_converter = RankingsConverter(config)
+        result = rankings_converter.convert(generic_dataset)
+
+        expected_result = {
+            "data": [
+                {
+                    "payload": [
+                        {
+                            "query": "query 1",
+                            "documents": ["passage 1", "passage 2"],
+                            "model": "test_model",
+                            "additional_key": "additional_value",
+                        }
+                    ]
+                },
+                {
+                    "payload": [
+                        {
+                            "query": "query 2",
+                            "documents": ["passage 3", "passage 4"],
+                            "model": "test_model",
+                            "additional_key": "additional_value",
+                        }
+                    ]
+                },
+            ]
+        }
+
+        assert result == expected_result
+
     @pytest.mark.parametrize(
         "queries_data, passages_data, expected_error",
         [