refactor

Author: mshr-h

AGENTS.md

@@ -40,6 +40,7 @@ onnx2fx/
 │       ├── analyze.py     # Model analysis utilities
 │       ├── attributes.py  # ONNX attribute parsing
 │       ├── dtype.py       # ONNX to PyTorch dtype mapping
+│       ├── external_data.py # External data utilities
 │       ├── names.py       # Name sanitization utilities
 │       ├── op_helpers.py  # Op helper utilities
 │       └── training.py    # Training utilities (make_trainable)
@@ -63,7 +64,6 @@ Core Requirements:
 - Python >= 3.11
 - PyTorch >= 2.9.0
 - ONNX >= 1.19.1
-- onnxscript >= 0.3.0
 
 Development Tools:
 - uv (recommended for development)
@@ -147,7 +147,7 @@ def bias_gelu(builder, node):
 - `builder.has_value(name)` - Check if value exists in environment
 - `builder.call_function(func, args, kwargs)` - Create function call node
 - `builder.call_module(module_name, args, kwargs)` - Create module call node
-- `builder.add_submodule(name, module)` - Register a submodule (returns safe name)
+- `builder.register_submodule(name, module)` - Register a submodule (returns safe name)
 - `builder.opset_version` - Get current opset version for default domain
 - `builder.get_opset_version(domain)` - Get opset version for specific domain
 
@@ -161,7 +161,7 @@ For parsing ONNX node attributes, use functions from `onnx2fx.utils.attributes`:
 ### Public API
 
 #### Core Functions
-- `convert(model)` - Convert ONNX model to FX GraphModule
+- `convert(model, *, base_dir=None, memmap_external_data=False)` - Convert ONNX model to FX GraphModule
 - `make_trainable(module)` - Convert buffers to trainable parameters for training
 
 #### Model Analysis

README.md

@@ -9,7 +9,7 @@ Yet another ONNX to PyTorch FX converter.
 ## Features
 
 - **Simple API**: Convert ONNX models with a single function call
-- **Extensive Operator Support**: 170+ ONNX operators including standard and Microsoft domain operators
+- **Extensive Operator Support**: Wide ONNX operator coverage including standard and Microsoft domain operators
 - **Multi-Opset Version Support**: Automatic selection of version-specific operator handlers based on model opset
 - **Custom Operator Registration**: Easily extend support for unsupported or custom ONNX operators
 - **PyTorch FX Output**: Get a `torch.fx.GraphModule` for easy inspection, optimization, and compilation
@@ -36,22 +36,21 @@ The following models have been tested and verified to work with onnx2fx:
 - Python >= 3.11
 - PyTorch >= 2.9.0
 - ONNX >= 1.19.1
-- onnxscript >= 0.3.0
 
 ### From Source
 
 ```bash
 git clone https://github.com/mshr-h/onnx2fx.git
 cd onnx2fx
-pip install .
+uv sync
 ```
 
 ### Development Installation
 
 ```bash
 git clone https://github.com/mshr-h/onnx2fx.git
 cd onnx2fx
-pip install -e ".[dev]"
+uv sync --dev
 ```
 
 ## Quick Start
@@ -70,6 +69,10 @@ fx_module = convert("model.onnx")
 onnx_model = onnx.load("model.onnx")
 fx_module = convert(onnx_model)
 
+# For models with external data, you can pass base_dir.
+# memmap_external_data avoids loading external data into memory.
+fx_module = convert("model.onnx", base_dir="/path/to/model_dir", memmap_external_data=True)
+
 # Run inference
 input_tensor = torch.randn(1, 3, 224, 224)
 output = fx_module(input_tensor)

pyproject.toml

@@ -9,7 +9,6 @@ authors = [
 requires-python = ">=3.11"
 dependencies = [
     "onnx>=1.19.1",
-    "onnxscript>=0.5.7",
     "torch>=2.9.0",
 ]
 
@@ -20,6 +19,7 @@ build-backend = "uv_build"
 [dependency-groups]
 dev = [
     "huggingface-hub>=0.36.0",
+    "onnxscript>=0.5.7",
     "onnxruntime>=1.23.2",
     "pre-commit>=4.5.1",
     "pytest>=9.0.2",