refractor(tests): make run arguments to model config

huggingface · Dec 4, 2024 · 3a5a7f0 · 3a5a7f0
1 parent 1af9edc
commit 3a5a7f0
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/text-generation-inference/integration-tests/conftest.py b/text-generation-inference/integration-tests/conftest.py
@@ -146,16 +146,6 @@ def docker_launcher(
         logger.info(f"Starting docker launcher for model {model_id}")
         port = 8080
 
-        args = [
-            "--max-input-length", "512",
-            "--max-total-tokens", "1024",
-            "--max-batch-prefill-tokens", "512",
-            "--max-batch-total-tokens", "1024"
-        ]
-
-        if trust_remote_code:
-            args.append("--trust-remote-code")
-
         client = docker.from_env()
 
         container_name = f"tgi-tests-{model_id.split('/')[-1]}"
@@ -171,6 +161,11 @@ def docker_launcher(
             logger.error(f"Error handling existing container: {str(e)}")
 
         model_name = next(name for name, cfg in MODEL_CONFIGS.items() if cfg["model_id"] == model_id)
+
+        args = MODEL_CONFIGS[model_name]["args"].copy()
+        if trust_remote_code:
+            args.append("--trust-remote-code")
+
         env = MODEL_CONFIGS[model_name]["env_config"].copy()
 
         # Add model_id to env

diff --git a/text-generation-inference/integration-tests/test_model.py b/text-generation-inference/integration-tests/test_model.py
@@ -12,6 +12,12 @@
         "expected_greedy_output": "\n\nDeep learning is a new field of research that has been around for a while",
         "expected_sampling_output": 'The fundamental concepts of deep learning are the same as those used to train and understand your first language, or your first set of skills',
         "expected_batch_output": "\n\nDeep learning is a technique that allows you to learn something from a single source",
+        "args": [
+            "--max-input-length", "512",
+            "--max-total-tokens", "1024",
+            "--max-batch-prefill-tokens", "512",
+            "--max-batch-total-tokens", "1024"
+        ],
         "env_config": {
             "LOG_LEVEL": "info,text_generation_router,text_generation_launcher=debug",
             "MAX_BATCH_SIZE": "4",
@@ -26,6 +32,12 @@
         "expected_greedy_output": "\n\nDeep learning is a subfield of machine learning that allows computers to learn from data",
         "expected_sampling_output": "Deep learning is a subfield of machine learning that focuses on mimicking the structure and function of the human brain",
         "expected_batch_output": "\n\nDeep learning is a subfield of machine learning that allows computers to learn from data",
+        "args": [
+            "--max-input-length", "512",
+            "--max-total-tokens", "1024",
+            "--max-batch-prefill-tokens", "512",
+            "--max-batch-total-tokens", "1024"
+        ],
         "env_config": {
             "LOG_LEVEL": "info,text_generation_router,text_generation_launcher=debug",
             "MAX_BATCH_SIZE": "4",