update models for TRT-LLM 0.6.1 (#1392)

deepjavalibrary · Dec 13, 2023 · 91c0e52 · 91c0e52
1 parent 81c538f
commit 91c0e52
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tests/integration/llm/prepare.py b/tests/integration/llm/prepare.py
@@ -626,12 +626,15 @@
         "option.output_formatter": "jsonlines",
     },
     "falcon-7b": {
-        "option.model_id": "s3://djl-llm/triton/falcon-7b-tp1-bs4/",
+        "option.model_id": "s3://djl-llm/triton/0.6.1/falcon-7b-tp1-bs16/",
         "option.tensor_parallel_degree": 1,
+        "option.max_input_len": 1024,
+        "option.max_output_len": 512,
+        "option.max_rolling_batch_size": 16,
         "option.output_formatter": "jsonlines",
     },
     "llama2-7b-smoothquant": {
-        "option.model_id": "TheBloke/Llama-2-7B-fp16",
+        "option.model_id": "s3://djl-llm/meta-llama-Llama-2-7b-chat-hf/",
         "option.tensor_parallel_degree": 4,
         "option.quantize": "smoothquant",
         "option.smoothquant_per_token": "True",