From bc5220edfb5e3744c2f77e6f5077f58d1ab0dc45 Mon Sep 17 00:00:00 2001
From: David Yastremsky <58150256+dyastremsky@users.noreply.github.com>
Date: Mon, 16 Dec 2024 14:39:20 -0800
Subject: [PATCH] Fix re-rankings payload schema for NIMs (#220)

---
 .../inputs/converters/rankings_converter.py   |  6 ++--
 genai-perf/tests/test_rankings_converter.py   | 36 +++++++++----------
 2 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/genai-perf/genai_perf/inputs/converters/rankings_converter.py b/genai-perf/genai_perf/inputs/converters/rankings_converter.py
index 7b807a59..f3613108 100644
--- 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
index ae5596b9..ad9cb7cc 100644
--- 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,8 +266,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",
                                 }
                             ]
@@ -275,8 +275,8 @@ def test_convert_missing_files(self, queries_data, passages_data, expected_error
                         {
                             "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",
                                 }
                             ]