skypro1111 commited on
Commit
e555bf0
·
verified ·
1 Parent(s): 01b0fdb

Update onnx/infer_onnx_hf.py

Browse files
Files changed (1) hide show
  1. onnx/infer_onnx_hf.py +2 -1
onnx/infer_onnx_hf.py CHANGED
@@ -8,7 +8,7 @@ from huggingface_hub import hf_hub_download
8
  model_name = "skypro1111/mbart-large-50-verbalization"
9
 
10
 
11
- def download_model_from_hf(repo_id=model_name):
12
  """Download ONNX models from HuggingFace Hub."""
13
 
14
  files = ["onnx/encoder_model.onnx", "onnx/decoder_model.onnx", "onnx/decoder_model.onnx_data"]
@@ -17,6 +17,7 @@ def download_model_from_hf(repo_id=model_name):
17
  hf_hub_download(
18
  repo_id=repo_id,
19
  filename=file,
 
20
  )
21
 
22
  return files
 
8
  model_name = "skypro1111/mbart-large-50-verbalization"
9
 
10
 
11
+ def download_model_from_hf(repo_id=model_name, model_dir="./"):
12
  """Download ONNX models from HuggingFace Hub."""
13
 
14
  files = ["onnx/encoder_model.onnx", "onnx/decoder_model.onnx", "onnx/decoder_model.onnx_data"]
 
17
  hf_hub_download(
18
  repo_id=repo_id,
19
  filename=file,
20
+ local_dir=model_dir,
21
  )
22
 
23
  return files