Fix re-rankings payload schema for NIMs (#220)

triton-inference-server · Dec 16, 2024 · bc5220e · bc5220e
1 parent 7db2604
commit bc5220e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 22 deletions.
diff --git a/genai-perf/genai_perf/inputs/converters/rankings_converter.py b/genai-perf/genai_perf/inputs/converters/rankings_converter.py
@@ -72,11 +72,9 @@ def convert(
                 passages = passage_entry.texts
                 payload = {"query": query, "texts": passages}
             else:
-                passages = [
-                    {"text_input": p} for p in passage_entry.texts if p is not None
-                ]
+                passages = [{"text": p} for p in passage_entry.texts if p is not None]
                 payload = {
-                    "query": query,
+                    "query": {"text": query},
                     "passages": passages,
                     "model": model_name,
                 }

diff --git a/genai-perf/tests/test_rankings_converter.py b/genai-perf/tests/test_rankings_converter.py
@@ -81,10 +81,10 @@ def test_convert_default(self):
                 {
                     "payload": [
                         {
-                            "query": "query 1",
+                            "query": {"text": "query 1"},
                             "passages": [
-                                {"text_input": "passage 1"},
-                                {"text_input": "passage 2"},
+                                {"text": "passage 1"},
+                                {"text": "passage 2"},
                             ],
                             "model": "test_model",
                         }
@@ -93,10 +93,10 @@ def test_convert_default(self):
                 {
                     "payload": [
                         {
-                            "query": "query 2",
+                            "query": {"text": "query 2"},
                             "passages": [
-                                {"text_input": "passage 3"},
-                                {"text_input": "passage 4"},
+                                {"text": "passage 3"},
+                                {"text": "passage 4"},
                             ],
                             "model": "test_model",
                         }
@@ -135,10 +135,10 @@ def test_convert_with_request_parameters(self):
                 {
                     "payload": [
                         {
-                            "query": "query 1",
+                            "query": {"text": "query 1"},
                             "passages": [
-                                {"text_input": "passage 1"},
-                                {"text_input": "passage 2"},
+                                {"text": "passage 1"},
+                                {"text": "passage 2"},
                             ],
                             "model": "test_model",
                             "encoding_format": "base64",
@@ -150,10 +150,10 @@ def test_convert_with_request_parameters(self):
                 {
                     "payload": [
                         {
-                            "query": "query 2",
+                            "query": {"text": "query 2"},
                             "passages": [
-                                {"text_input": "passage 3"},
-                                {"text_input": "passage 4"},
+                                {"text": "passage 3"},
+                                {"text": "passage 4"},
                             ],
                             "model": "test_model",
                             "encoding_format": "base64",
@@ -266,17 +266,17 @@ def test_convert_missing_files(self, queries_data, passages_data, expected_error
                         {
                             "payload": [
                                 {
-                                    "query": "query 1",
-                                    "passages": [{"text_input": "passage 1"}],
+                                    "query": {"text": "query 1"},
+                                    "passages": [{"text": "passage 1"}],
                                     "model": "test_model",
                                 }
                             ]
                         },
                         {
                             "payload": [
                                 {
-                                    "query": "query 2",
-                                    "passages": [{"text_input": "passage 2"}],
+                                    "query": {"text": "query 2"},
+                                    "passages": [{"text": "passage 2"}],
                                     "model": "test_model",
                                 }
                             ]
@@ -293,8 +293,8 @@ def test_convert_missing_files(self, queries_data, passages_data, expected_error
                         {
                             "payload": [
                                 {
-                                    "query": "query 1",
-                                    "passages": [{"text_input": "passage 1"}],
+                                    "query": {"text": "query 1"},
+                                    "passages": [{"text": "passage 1"}],
                                     "model": "test_model",
                                 }
                             ]