File size: 2,073 Bytes
9f3aa4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | import argparse
import os
from pathlib import Path
from dotenv import load_dotenv
from huggingface_hub import HfApi
load_dotenv()
MODELS_ROOT = Path("models")
VALID_MODES = ("marker", "qa_m", "qa_b", "fasttext")
REPO_PREFIX = "lamossta/distillbert"
FASTTEXT_REPO = "lamossta/fasttext_baseline"
def _repo_id(mode: str) -> str:
if mode == "fasttext":
return FASTTEXT_REPO
return f"{REPO_PREFIX}_{mode}"
def _model_filename(mode: str) -> str:
return "model.bin" if mode == "fasttext" else "model.onnx"
def upload_model(mode: str, revision: str = "main") -> None:
token = os.environ.get("HF_TOKEN")
api = HfApi(token=token)
model_dir = MODELS_ROOT / mode
repo_id = _repo_id(mode)
filename = _model_filename(mode)
model_path = model_dir / filename
if not model_path.exists():
raise FileNotFoundError(f"Model file not found: {model_path}")
api.create_repo(repo_id, exist_ok=True, token=token)
api.upload_file(
repo_id=repo_id,
path_or_fileobj=str(model_path),
path_in_repo=filename,
revision=revision,
token=token,
)
print(f"Uploaded '{mode}' model to {repo_id}/{filename}")
def upload_all(revision: str = "main") -> list[str]:
uploaded = []
for mode in VALID_MODES:
model_path = MODELS_ROOT / mode / _model_filename(mode)
if not model_path.exists():
print(f"Skipping '{mode}': {model_path} not found")
continue
upload_model(mode, revision)
uploaded.append(mode)
return uploaded
def main():
parser = argparse.ArgumentParser(description="Upload ONNX models to Hugging Face")
parser.add_argument("--mode", default=None, choices=VALID_MODES, help="Single mode to upload (default: all)")
parser.add_argument("--revision", default="main")
args = parser.parse_args()
if args.mode:
upload_model(args.mode, args.revision)
else:
uploaded = upload_all(args.revision)
print(f"Uploaded models: {uploaded}")
if __name__ == "__main__":
main()
|