remove_weights_from_python_wheel

MANIFEST.in

@@ -1,6 +1,7 @@
 include README.md
 include THIRD_PARTY_NOTICES.md
-recursive-include nemotron_table_structure_v1
+recursive-include nemotron_table_structure_v1 *.py *.json
+recursive-exclude nemotron_table_structure_v1 *.pth
 
 
 

README.md

@@ -148,8 +148,12 @@ import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image
 
-from model import define_model
-from utils import plot_sample, postprocess_preds_table_structure, reformat_for_plotting
+from nemotron_table_structure_v1 import (
+    define_model,
+    plot_sample,
+    postprocess_preds_table_structure,
+    reformat_for_plotting,
+)
 
 # Load image
 path = "./example.png"

nemotron_table_structure_v1/__init__.py

@@ -9,7 +9,7 @@ A specialized object detection model for table structure extraction based on YOL
 
 __version__ = "1.0.0"
 
-from .model import define_model, YoloXWrapper
+from .model import define_model, YoloXWrapper, get_weights_path
 from .utils import (
     plot_sample,
     postprocess_preds_table_structure,
@@ -19,6 +19,7 @@ from .utils import (
 
 __all__ = [
     "define_model",
+    "get_weights_path",
     "YoloXWrapper",
     "plot_sample",
     "postprocess_preds_table_structure",

nemotron_table_structure_v1/model.py

@@ -10,8 +10,37 @@ import numpy.typing as npt
 import torch.nn as nn
 import torch.nn.functional as F
 from typing import Dict, List, Tuple, Union
+from huggingface_hub import hf_hub_download
 from .yolox.boxes import postprocess
 
+# HuggingFace repository for weights
+HF_REPO_ID = "nvidia/nemotron-table-structure-v1"
+WEIGHTS_FILENAME = "weights.pth"
+
+
+def get_weights_path(verbose: bool = True) -> str:
+    """
+    Get the path to the model weights, downloading from HuggingFace if necessary.
+
+    The weights are cached in the HuggingFace cache directory after the first download.
+
+    Args:
+        verbose (bool): Whether to print download progress. Defaults to True.
+
+    Returns:
+        str: Path to the weights file.
+    """
+    if verbose:
+        print(f" -> Downloading/loading weights from HuggingFace: {HF_REPO_ID}")
+
+    weights_path = hf_hub_download(
+        repo_id=HF_REPO_ID,
+        filename=WEIGHTS_FILENAME,
+        repo_type="model",
+    )
+
+    return weights_path
+
 
 def define_model(config_name: str = "page_element_v3", verbose: bool = True) -> nn.Module:
     """
@@ -32,13 +61,8 @@ def define_model(config_name: str = "page_element_v3", verbose: bool = True) ->
     config = exp_module.Exp()
     model = config.get_model()
 
-    # Load weights
-    if verbose:
-        print(" -> Loading weights from", config.ckpt)
-
-    # Find package directory and load weights (nemotron_table_structure_v1)
-    package_dir = os.path.dirname(os.path.abspath(__file__))
-    weights_path = os.path.join(package_dir, "weights.pth")
+    # Load weights (downloaded from HuggingFace if not cached)
+    weights_path = get_weights_path(verbose=verbose)
     state_dict = torch.load(weights_path, map_location="cpu", weights_only=False)
     model.load_state_dict(state_dict["model"], strict=True)
 

pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
     "matplotlib>=3.3.0",
     "pandas>=1.3.0",
     "Pillow>=8.0.0",
+    "huggingface_hub>=0.20.0",
 ]
 
 [project.urls]
@@ -45,7 +46,7 @@ Documentation = "https://huggingface.co/nvidia/nemotron-table-structure-v1"
 packages = ["nemotron_table_structure_v1", "nemotron_table_structure_v1.yolox", "nemotron_table_structure_v1.post_processing"]
 
 [tool.setuptools.package-data]
-"*" = ["*.pth", "config.json"]
+"*" = ["config.json"]