hf upload/download and onnx export

src/models/hf_download.py

@@ -0,0 +1,68 @@
+import argparse
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+from huggingface_hub import hf_hub_download
+
+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 download_model(mode: str, revision: str = "main") -> Path:
+    token = os.environ.get("HF_TOKEN")
+    repo_id = _repo_id(mode)
+    model_dir = MODELS_ROOT / mode
+    model_dir.mkdir(parents=True, exist_ok=True)
+    filename = _model_filename(mode)
+
+    hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        revision=revision,
+        token=token,
+        local_dir=str(model_dir),
+    )
+    print(f"Downloaded {repo_id}/{filename} -> {model_dir / filename}")
+
+    return model_dir
+
+
+def download_all(revision: str = "main") -> dict[str, Path]:
+    downloaded = {}
+    for mode in VALID_MODES:
+        try:
+            downloaded[mode] = download_model(mode, revision)
+        except Exception as e:
+            print(f"Skipping '{mode}': {e}")
+    return downloaded
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Download ONNX models from Hugging Face")
+    parser.add_argument("--mode", default=None, choices=VALID_MODES, help="Single mode to download (default: all)")
+    parser.add_argument("--revision", default="main")
+    args = parser.parse_args()
+
+    if args.mode:
+        download_model(args.mode, args.revision)
+    else:
+        downloaded = download_all(args.revision)
+        print(f"Downloaded models: {downloaded}")
+
+
+if __name__ == "__main__":
+    main()

src/models/hf_upload.py

@@ -0,0 +1,74 @@
+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()

src/models/serialize.py

@@ -0,0 +1,63 @@
+import argparse
+from pathlib import Path
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+
+def export_to_onnx(
+    model_dir: str,
+    output_path: str | None = None,
+    max_len: int = 256,
+    opset_version: int = 17,
+) -> Path:
+    model_dir = Path(model_dir)
+    if output_path is None:
+        output_path = model_dir / "model.onnx"
+    else:
+        output_path = Path(output_path)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+    model.eval()
+
+    dummy = tokenizer(
+        ["dummy input", "another dummy"],
+        max_length=max_len,
+        truncation=True,
+        padding="max_length",
+        return_tensors="pt",
+    )
+
+    batch = torch.export.Dim("batch", min=1, max=4096)
+    torch.onnx.export(
+        model,
+        (dummy["input_ids"], dummy["attention_mask"]),
+        str(output_path),
+        input_names=["input_ids", "attention_mask"],
+        output_names=["logits"],
+        dynamic_shapes={
+            "input_ids": {0: batch},
+            "attention_mask": {0: batch},
+        },
+        opset_version=opset_version,
+        dynamo=True,
+        external_data=False,
+    )
+
+    print(f"Exported ONNX model to '{output_path}' ({output_path.stat().st_size / 1024 / 1024:.1f} MB)")
+    return output_path
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Export model to ONNX")
+    parser.add_argument("--model-dir", required=True, help="Path to saved model directory")
+    parser.add_argument("--output", default=None, help="Output ONNX path (default: model_dir/model.onnx)")
+    parser.add_argument("--max-len", type=int, default=256)
+    parser.add_argument("--opset", type=int, default=17)
+    args = parser.parse_args()
+
+    export_to_onnx(args.model_dir, args.output, args.max_len, args.opset)
+
+
+if __name__ == "__main__":
+    main()