badaoui HF Staff commited on
Commit
6f698cf
·
verified ·
1 Parent(s): c7f5d45

bigger bs, seq_len for text-generation task

Browse files
Files changed (1) hide show
  1. 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": 1, "sequence_length": 128, "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":
 
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":