Update onnx/infer_onnx_hf.py
Browse files- 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
|