Spaces:
Running
Running
bigger bs, seq_len for text-generation task
Browse files- optimum_neuron_export.py +1 -1
optimum_neuron_export.py
CHANGED
|
@@ -151,7 +151,7 @@ def get_default_inputs(task_or_pipeline: str, pipeline_name: str = None) -> Dict
|
|
| 151 |
if task_or_pipeline in ["feature-extraction", "sentence-transformers", "fill-mask", "question-answering", "text-classification", "token-classification"]:
|
| 152 |
return {"batch_size": 1, "sequence_length": 128}
|
| 153 |
elif task_or_pipeline == "text_generation":
|
| 154 |
-
return {"batch_size":
|
| 155 |
elif task_or_pipeline == "multiple-choice":
|
| 156 |
return {"batch_size": 1, "num_choices": 4, "sequence_length": 128}
|
| 157 |
elif task_or_pipeline == "text2text-generation":
|
|
|
|
| 151 |
if task_or_pipeline in ["feature-extraction", "sentence-transformers", "fill-mask", "question-answering", "text-classification", "token-classification"]:
|
| 152 |
return {"batch_size": 1, "sequence_length": 128}
|
| 153 |
elif task_or_pipeline == "text_generation":
|
| 154 |
+
return {"batch_size": 4, "sequence_length": 4096, "tensor_parallel_size": 4}
|
| 155 |
elif task_or_pipeline == "multiple-choice":
|
| 156 |
return {"batch_size": 1, "num_choices": 4, "sequence_length": 128}
|
| 157 |
elif task_or_pipeline == "text2text-generation":
|