From 229e5e85eeef481916d8bf67f24bcf9dfcb68b25 Mon Sep 17 00:00:00 2001
From: Sai Kiran Polisetty <spolisetty@nvidia.com>
Date: Mon, 30 Sep 2024 12:04:58 +0530
Subject: [PATCH] Update CI

---
 qa/L0_http/generate_endpoint_test.py          | 21 ++++++-------------
 .../generate_models/mock_llm/config.pbtxt     |  8 +------
 2 files changed, 7 insertions(+), 22 deletions(-)

diff --git a/qa/L0_http/generate_endpoint_test.py b/qa/L0_http/generate_endpoint_test.py
index 3eb0b6ea5f..7cd6331aa2 100755
--- a/qa/L0_http/generate_endpoint_test.py
+++ b/qa/L0_http/generate_endpoint_test.py
@@ -144,8 +144,7 @@ def test_generate(self):
 
     def test_generate_with_all_inputs(self):
         # Setup text-based input
-        text = "hello world"
-        inputs = {"PROMPT": text, "STREAM": False, "input_ids": [100, 200]}
+        inputs = {"PROMPT": ["hello", "world"], "STREAM": False}
 
         r = self.generate(self._model_name, inputs)
         r.raise_for_status()
@@ -155,7 +154,7 @@ def test_generate_with_all_inputs(self):
 
         data = r.json()
         self.assertIn("TEXT", data)
-        self.assertEqual(text, data["TEXT"])
+        self.assertEqual(["hello", "world"], data["TEXT"])
 
     def test_request_id(self):
         # Setup text based input
@@ -235,26 +234,18 @@ def test_missing_inputs(self):
         ]
         for inputs in missing_all_inputs:
             self.generate_expect_failure(
-                self._model_name,
-                inputs,
-                "expected number of inputs between 2 and 3 but got 0",
+                self._model_name, inputs, "expected 2 inputs but got 0"
             )
             self.generate_stream_expect_failure(
-                self._model_name,
-                inputs,
-                "expected number of inputs between 2 and 3 but got 0",
+                self._model_name, inputs, "expected 2 inputs but got 0"
             )
 
         for inputs in missing_one_input:
             self.generate_expect_failure(
-                self._model_name,
-                inputs,
-                "expected number of inputs between 2 and 3 but got 1",
+                self._model_name, inputs, "expected 2 inputs but got 1"
             )
             self.generate_stream_expect_failure(
-                self._model_name,
-                inputs,
-                "expected number of inputs between 2 and 3 but got 1",
+                self._model_name, inputs, "expected 2 inputs but got 1"
             )
 
     def test_invalid_input_types(self):
diff --git a/qa/L0_http/generate_models/mock_llm/config.pbtxt b/qa/L0_http/generate_models/mock_llm/config.pbtxt
index 74a306052a..424049f3f9 100644
--- a/qa/L0_http/generate_models/mock_llm/config.pbtxt
+++ b/qa/L0_http/generate_models/mock_llm/config.pbtxt
@@ -35,18 +35,12 @@ input [
   {
     name: "PROMPT"
     data_type: TYPE_STRING
-    dims: [ 1, 1 ]
+    dims: [ 1, -1 ]
   },
   {
     name: "STREAM"
     data_type: TYPE_BOOL
     dims: [ 1, 1 ]
-  },
-  {
-    name: "input_ids"
-    data_type: TYPE_INT32
-    dims: [ 1, -1 ]
-    optional: true
   }
 ]