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