[Enhancement] Pass model_init_kwargs to sparse model auto tracing to load remote model class (#555)

* add trust remote code

Signed-off-by: zhichao-aws <[email protected]>

* pass model kwargs

Signed-off-by: zhichao-aws <[email protected]>

* sanitize, lint

Signed-off-by: zhichao-aws <[email protected]>

* add ut

Signed-off-by: zhichao-aws <[email protected]>

* changelog

Signed-off-by: zhichao-aws <[email protected]>

* fix json string in GH action

Signed-off-by: zhichao-aws <[email protected]>

---------

Signed-off-by: zhichao-aws <[email protected]>

Author: zhichao-aws

.ci/run-repository.sh

@@ -77,20 +77,21 @@ elif [[ "$TASK_TYPE" == "SentenceTransformerTrace" || "$TASK_TYPE" == "SparseTra
   echo -e "\033[34;1mINFO:\033[0m ACTIVATION: ${ACTIVATION:-N/A}\033[0m"
   echo -e "\033[34;1mINFO:\033[0m MODEL_DESCRIPTION: ${MODEL_DESCRIPTION:-N/A}\033[0m"
   echo -e "\033[34;1mINFO:\033[0m MODEL_NAME: ${MODEL_NAME:-N/A}\033[0m"
+  echo -e "\033[34;1mINFO:\033[0m MODEL_INIT_KWARGS: ${MODEL_INIT_KWARGS:-{}}\033[0m"
 
   if [[ "$TASK_TYPE" == "SentenceTransformerTrace" ]]; then
       NOX_TRACE_TYPE="trace"
-      EXTRA_ARGS="-ed ${EMBEDDING_DIMENSION} -pm ${POOLING_MODE}"
+      EXTRA_ARGS=( -ed "${EMBEDDING_DIMENSION}" -pm "${POOLING_MODE}" )
   elif [[ "$TASK_TYPE" == "SparseTrace" ]]; then
       NOX_TRACE_TYPE="sparsetrace"
-      EXTRA_ARGS="-spr ${SPARSE_PRUNE_RATIO} -act ${ACTIVATION}"
+      EXTRA_ARGS=( -spr "${SPARSE_PRUNE_RATIO}" -act "${ACTIVATION}" -mik "${MODEL_INIT_KWARGS}" )
   elif [[ "$TASK_TYPE" == "SparseTokenizerTrace" ]]; then
       NOX_TRACE_TYPE="sparsetrace"
-      # use extra args to trigger the tokenizer tracing logics
-      EXTRA_ARGS="-t"
+      # use extra args to trigger the tokenizer tracing logics (no -mik for tokenizer)
+      EXTRA_ARGS=( -t )
   elif [[ "$TASK_TYPE" == "SemanticHighlighterTrace" ]]; then
       NOX_TRACE_TYPE="semantic_highlighter_trace"
-      EXTRA_ARGS=""
+      EXTRA_ARGS=()
   else
       echo "Unknown TASK_TYPE: $TASK_TYPE"
       exit 1
@@ -105,7 +106,7 @@ elif [[ "$TASK_TYPE" == "SentenceTransformerTrace" || "$TASK_TYPE" == "SparseTra
     -up "${UPLOAD_PREFIX}"
     -mn "${MODEL_NAME}"
     -md "${MODEL_DESCRIPTION:+"$MODEL_DESCRIPTION"}"
-    ${EXTRA_ARGS}
+    "${EXTRA_ARGS[@]}"
   )
 
   echo "nox -s ${nox_command[@]}"

.github/workflows/model_uploader.yml

@@ -67,6 +67,7 @@ on:
             "sparse_prune_ratio": (Optional) Float. Specifies the model-side prune ratio based on max values. Sparse model only.
             "activation": (Optional) String. Specifies the activation function for the sparse model. Sparse model only.
             "model_name": (Optional) String. Specifies the model name for uploading. Example: transforms "sentence-transformers/model" to "sentence-transformers/{model_name}",
+            "model_init_kwargs": (Optional) Object. JSON object to pass to from_pretrained via **kwargs.
           }
           
           Example:
@@ -102,23 +103,27 @@ jobs:
         sparse_prune_ratio=0
         activation=""
         model_name="${model_id##*/}"
+        model_init_kwargs="{}"
 
         if [ "$custom_params" != "{}" ] && [ -n "$custom_params" ]; then
           tmp_up=$(echo "$custom_params" | jq -r '.upload_prefix | select(.!=null)')
           tmp_spr=$(echo "$custom_params" | jq -r '.sparse_prune_ratio | select(.!=null)')
           tmp_act=$(echo "$custom_params" | jq -r '.activation | select(.!=null)')
           tmp_mn=$(echo "$custom_params" | jq -r '.model_name | select(.!=null)')
+          tmp_mik=$(echo "$custom_params" | jq -c '.model_init_kwargs | select(.!=null)')
 
           [ -n "$tmp_up" ] && upload_prefix="$tmp_up"
           [ -n "$tmp_spr" ] && sparse_prune_ratio="$tmp_spr"
           [ -n "$tmp_act" ] && activation="$tmp_act"
           [ -n "$tmp_mn" ] && model_name="$tmp_mn"
+          [ -n "$tmp_mik" ] && model_init_kwargs="$tmp_mik"
         fi
 
         echo "upload_prefix=$upload_prefix" >> $GITHUB_OUTPUT
         echo "sparse_prune_ratio=$sparse_prune_ratio" >> $GITHUB_OUTPUT
         echo "activation=$activation" >> $GITHUB_OUTPUT
         echo "model_name=$model_name" >> $GITHUB_OUTPUT
+        echo "model_init_kwargs=$model_init_kwargs" >> $GITHUB_OUTPUT
     - name: Initiate folders
       # This scripts init the folders path variables.
       # 1. Retrieves the input model_id.
@@ -167,6 +172,7 @@ jobs:
         - Model Prefix Folder: ${{ steps.init_folders.outputs.model_prefix_folder }}
         - Sparse Prune Ratio: ${{ steps.parse_custom_params.outputs.sparse_prune_ratio || 'N/A' }}
         - Activation: ${{ steps.parse_custom_params.outputs.activation || 'N/A' }}
+        - Model Init Kwargs: '${{ toJSON(fromJSON(steps.parse_custom_params.outputs.model_init_kwargs || '{}')) }}'
         
         ======== Workflow Output Information =========
         - Embedding Verification: Passed"
@@ -190,6 +196,7 @@ jobs:
       sparse_prune_ratio: ${{ steps.parse_custom_params.outputs.sparse_prune_ratio }}
       activation: ${{ steps.parse_custom_params.outputs.activation }}
       model_name: ${{ steps.parse_custom_params.outputs.model_name }}
+      model_init_kwargs: ${{ steps.parse_custom_params.outputs.model_init_kwargs }}
 
   # Step 3: Check if the model already exists in the model hub
   checking-out-model-hub:
@@ -267,6 +274,9 @@ jobs:
           echo "MODEL_NAME=${{ needs.init-workflow-var.outputs.model_name }}" >> $GITHUB_ENV
           echo "SPARSE_PRUNE_RATIO=${{ needs.init-workflow-var.outputs.sparse_prune_ratio }}" >> $GITHUB_ENV
           echo "ACTIVATION=${{ needs.init-workflow-var.outputs.activation }}" >> $GITHUB_ENV
+          echo "MODEL_INIT_KWARGS<<EOF" >> $GITHUB_ENV
+          echo '${{ toJSON(fromJSON(needs.init-workflow-var.outputs.model_init_kwargs)) }}' >> $GITHUB_ENV
+          echo "EOF" >> $GITHUB_ENV
       - name: Autotracing ${{ matrix.cluster }} secured=${{ matrix.secured }} version=${{matrix.entry.opensearch_version}}
         run: "./.ci/run-tests ${{ matrix.cluster }} ${{ matrix.secured }} ${{ matrix.entry.opensearch_version }} ${{github.event.inputs.model_type}}Trace"
         shell: bash

CHANGELOG.md

@@ -6,6 +6,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ### Added
 - Add space type mapping for sentence transformer models by @nathaliellenaa in ([#512](https://github.com/opensearch-project/opensearch-py-ml/pull/512))
 - Add example script for deploying semantic highlighter model on aws sagemaker. ([#513](https://github.com/opensearch-project/opensearch-py-ml/pull/513))
+- Add model_init_kwargs to sparse model uploading pipeline. ([#555](https://github.com/opensearch-project/opensearch-py-ml/pull/555))
 
 ### Changed
 - Update model upload history -  opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill (v.1.0.0)(TORCH_SCRIPT) by @dhrubo-os ([#415](https://github.com/opensearch-project/opensearch-py-ml/pull/415))

opensearch_py_ml/ml_models/sparse_encoding_model.py

@@ -6,6 +6,8 @@
 # GitHub history for details.
 import json
 import os
+import re
+from typing import Optional
 from zipfile import ZipFile
 
 import torch
@@ -36,6 +38,30 @@ def _generate_default_model_description() -> str:
     return description
 
 
+def _sanitize_module_name(name: str) -> str:
+    name = re.sub(r"[^0-9A-Za-z_\.]", "_", name)
+    parts = []
+    for p in name.split("."):
+        if not p:
+            continue
+        if p[0].isdigit():
+            p = f"n_{p}"
+        parts.append(p)
+    return ".".join(parts)
+
+
+def sanitize_model_modules(model: torch.nn.Module) -> None:
+    seen: set[type] = set()
+    for m in model.modules():
+        cls = m.__class__
+        if cls in seen:
+            continue
+        safe = _sanitize_module_name(getattr(cls, "__module__", ""))
+        if safe and safe != cls.__module__:
+            cls.__module__ = safe
+        seen.add(cls)
+
+
 class SparseEncodingModel(SparseModel):
     """
     Class for  exporting and configuring the NeuralSparseV2Model model.
@@ -50,12 +76,15 @@ def __init__(
         overwrite: bool = False,
         sparse_prune_ratio: float = 0,
         activation: str = None,
+        model_init_kwargs: Optional[dict] = None,
     ) -> None:
 
         super().__init__(model_id, folder_path, overwrite)
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        if model_init_kwargs is None:
+            model_init_kwargs = {}
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id, **model_init_kwargs)
         self.backbone_model = AutoModelForMaskedLM.from_pretrained(
-            model_id, _attn_implementation="eager"
+            model_id, _attn_implementation="eager", **model_init_kwargs
         )
         default_folder_path = os.path.join(
             os.getcwd(), "opensearch_neural_sparse_model_files"
@@ -167,6 +196,7 @@ def save_as_pt(
             return_tensors="pt",
         ).to(device)
 
+        sanitize_model_modules(cpu_model)
         compiled_model = torch.jit.trace(cpu_model, dict(features), strict=False)
         torch.jit.save(compiled_model, model_path)
         print("model file is saved to ", model_path)

tests/ml_models/test_sparseencondingmodel_pytest.py

@@ -13,8 +13,11 @@
 from zipfile import ZipFile
 
 import pytest
+from torch import nn
 
 from opensearch_py_ml.ml_models import SparseEncodingModel
+from opensearch_py_ml.ml_models.sparse_encoding_model import sanitize_model_modules
+from utils.model_uploader.autotracing_utils import init_sparse_model
 
 TEST_FOLDER = os.path.join(
     os.path.dirname(os.path.abspath("__file__")), "tests", "test_model_files"
@@ -374,5 +377,69 @@ def test_process_sparse_encoding():
     check_value(1.0706572532653809, encoding_result[1]["hello"], 0.001)
 
 
+def test_sanitize_module_name_and_trace():
+    class WeirdSub(nn.Module):
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, input_ids=None, attention_mask=None):
+            pass
+
+    # simulate weird remote module path
+    WeirdSub.__module__ = "remote.repo@bad:name/1"
+
+    class Toy(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.m = WeirdSub()
+
+        def forward(self, features: dict):
+            pass
+
+    model = Toy().eval()
+    sanitize_model_modules(model)
+
+    # After sanitize, module name should only contain [0-9A-Za-z_.]
+    assert all(c.isalnum() or c in {"_", "."} for c in model.m.__class__.__module__)
+
+
+def test_init_sparse_model_kwargs_passthrough():
+    received = {}
+
+    class FakeModel:
+        def __init__(
+            self,
+            model_id,
+            folder_path,
+            overwrite,
+            sparse_prune_ratio,
+            activation,
+            model_init_kwargs,
+        ):
+            received["model_id"] = model_id
+            received["folder_path"] = folder_path
+            received["overwrite"] = overwrite
+            received["sparse_prune_ratio"] = sparse_prune_ratio
+            received["activation"] = activation
+            received["model_init_kwargs"] = model_init_kwargs
+
+    model = init_sparse_model(
+        FakeModel,
+        model_id="foo/bar",
+        folder_path="/tmp/xyz",
+        sparse_prune_ratio=0.2,
+        activation="l0",
+        model_init_kwargs={"trust_remote_code": True, "revision": "dev"},
+    )
+
+    assert isinstance(model, FakeModel)
+    assert received["model_id"] == "foo/bar"
+    assert received["folder_path"] == "/tmp/xyz"
+    assert received["overwrite"] is True
+    assert received["sparse_prune_ratio"] == 0.2
+    assert received["activation"] == "l0"
+    assert received["model_init_kwargs"]["trust_remote_code"] is True
+
+
 clean_test_folder(TEST_FOLDER)
 clean_test_folder(TESTDATA_UNZIP_FOLDER)

utils/model_uploader/autotracing_utils.py

@@ -8,7 +8,7 @@
 import os
 import shutil
 import warnings
-from typing import Type, TypeVar
+from typing import Any, Dict, Optional, Type, TypeVar
 
 from huggingface_hub import HfApi
 
@@ -230,15 +230,23 @@ def __init__(self, stage: str, model_format: str, original_exception: Exception)
 
 
 def init_sparse_model(
-    model_class: Type[T], model_id, folder_path, sparse_prune_ratio=0, activation=None
+    model_class: Type[T],
+    model_id,
+    folder_path,
+    sparse_prune_ratio=0,
+    activation=None,
+    model_init_kwargs: Optional[Dict[str, Any]] = None,
 ) -> T:
     try:
+        if model_init_kwargs is None:
+            model_init_kwargs = {}
         pre_trained_model = model_class(
             model_id=model_id,
             folder_path=folder_path,
             overwrite=True,
             sparse_prune_ratio=sparse_prune_ratio,
             activation=activation,
+            model_init_kwargs=model_init_kwargs,
         )
     except Exception as e:
         raise ModelTraceError("initiating a sparse encoding model class object", e)