Fixing KeyError in tests related to default name change

PAIR-code · Oct 11, 2024 · 64114d5 · 64114d5
1 parent ae5825a
commit 64114d5
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/lit_nlp/examples/gcp/model_server_test.py b/lit_nlp/examples/gcp/model_server_test.py
@@ -25,9 +25,9 @@ def test_predict_endpoint(self, mock_get_models):
     ]
 
     mock_get_models.return_value = {
-        'gemma_1.1_instruct_2b_en': mock_model,
-        'gemma_1.1_instruct_2b_en_salience': salience_model,
-        'gemma_1.1_instruct_2b_en_tokenize': tokenize_model,
+        'gemma_1.1_2b_IT': mock_model,
+        'gemma_1.1_2b_IT_salience': salience_model,
+        'gemma_1.1_2b_IT_tokenize': tokenize_model,
     }
     app = webtest.TestApp(model_server.get_wsgi_app())