feat: add Cohere API support to rankings converter

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

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",
         [