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/model.py

@@ -10,35 +10,45 @@ 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 downloading model weights
+HF_REPO_ID = "nvidia/nemotron-table-structure-v1"
+WEIGHTS_FILENAME = "nemotron_table_structure_v1/weights.pth"
 
-def define_model(config_name: str = "page_element_v3", verbose: bool = True) -> nn.Module:
+
+def define_model(config_name: str = "table_structure_v1", verbose: bool = True) -> nn.Module:
     """
     Defines and initializes the model based on the configuration.
 
     Args:
-        config_name (str): Configuration name. Defaults to "page_element_v3".
+        config_name (str): Configuration name. Defaults to "table_structure_v1".
         verbose (bool): Whether to print verbose output. Defaults to True.
 
     Returns:
         torch.nn.Module: The initialized YOLOX model.
     """
     # Load model from exp_file
-    # page_element_v3.py is in the same directory as model.py
+    # table_structure_v1.py is in the same directory as model.py
     sys.path.append(os.path.dirname(__file__))
     exp_module = importlib.import_module("table_structure_v1")
 
     config = exp_module.Exp()
     model = config.get_model()
 
-    # Load weights
+    # Download weights from HuggingFace Hub (cached locally after first download)
+    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,
+    )
+
     if verbose:
-        print(" -> Loading weights from", config.ckpt)
+        print(f" -> Weights cached at: {weights_path}")
 
-    # 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")
     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"]