Modify the output to be a list to ensure compatibility with JSON seri…

…alization.
PAIR-code · Sep 5, 2024 · c9c7aff · c9c7aff
1 parent 9d39a00
commit c9c7aff
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/lit_nlp/examples/gcp/model_server.py b/lit_nlp/examples/gcp/model_server.py
@@ -27,7 +27,7 @@ def _handler(app, request, environ):
       data = serialize.from_json(request.data) if len(request.data) else None
       inputs = data['inputs']
       outputs = predict_fn(inputs)
-      response_body = serialize.to_json(outputs,  simple=True)
+      response_body = serialize.to_json(list(outputs),  simple=True)
       return app.respond(request, response_body, 'application/json', 200)
 
     return _handler

diff --git a/lit_nlp/examples/gcp/model_server_test.py b/lit_nlp/examples/gcp/model_server_test.py
@@ -10,18 +10,18 @@ class TestWSGIApp(absltest.TestCase):
   def test_predict_endpoint(self, mock_get_models):
 
     mock_model = mock.MagicMock()
-    mock_model.predict.side_effect = [{'response': 'test output text'}]
+    mock_model.predict.side_effect = [[{'response': 'test output text'}]]
 
     salience_model = mock.MagicMock()
-    salience_model.predict.side_effect = [{
+    salience_model.predict.side_effect = [[{
         'tokens': ['test', 'output', 'text'],
         'grad_l2': [0.1234, 0.3456, 0.5678],
         'grad_dot_input': [0.1234, -0.3456, 0.5678],
-    }]
+    }]]
 
     tokenize_model = mock.MagicMock()
     tokenize_model.predict.side_effect = [
-        {'tokens': ['test', 'output', 'text']}
+        [{'tokens': ['test', 'output', 'text']}]
     ]
 
     mock_get_models.return_value = {
@@ -33,22 +33,22 @@ def test_predict_endpoint(self, mock_get_models):
 
     response = app.post_json('/predict', {'inputs': 'test_input'})
     self.assertEqual(response.status_code, 200)
-    self.assertEqual(response.json, {'response': 'test output text'})
+    self.assertEqual(response.json, [{'response': 'test output text'}])
 
     response = app.post_json('/salience', {'inputs': 'test_input'})
     self.assertEqual(response.status_code, 200)
     self.assertEqual(
         response.json,
-        {
+        [{
             'tokens': ['test', 'output', 'text'],
             'grad_l2': [0.1234, 0.3456, 0.5678],
             'grad_dot_input': [0.1234, -0.3456, 0.5678],
-        },
+        }],
     )
 
     response = app.post_json('/tokenize', {'inputs': 'test_input'})
     self.assertEqual(response.status_code, 200)
-    self.assertEqual(response.json, {'tokens': ['test', 'output', 'text']})
+    self.assertEqual(response.json, [{'tokens': ['test', 'output', 'text']}])
 
 
 if __name__ == '__main__':