Custom model registration

View / edit source ↗

简体中文 | Documentation | Training | Historical registry tutorial

Registration connects a Python implementation to a configuration name. It does not download weights, make an arbitrary Transformers model compatible, provide training/export support, or certify model quality. This guide describes the checked-in registry, AutoModel, hub loader and dynamic importer.

Names and Import Order

Use @tables.register("model_classes", "YourUniqueModelName") on a class imported before construction. tables comes from funasr.register. The second argument is the exact, case-sensitive registry key; omitting it uses the Python class name. The decorator returns the original class and records its source location with inspect, so define examples in importable .py files, not only a REPL or dynamically generated class.

Keys live in process-global dictionaries. An existing key is overwritten, with a debug log rather than an error; imports and their order therefore matter. Use an organization/project-specific name and a collision guard. Avoid names such as SenseVoiceSmall, Paraformer or FunASRNano for unrelated custom models. tables.print("model") displays registration metadata; tables.model_classes[name] gives the active implementation.

Other tables include encoder_classes, decoder_classes, frontend_classes, tokenizer_classes, dataset_classes, index_ds_classes and batch_sampler_classes. Each consumer has its own contract: registering an encoder does not make it a complete AutoModel model. The registry also permits new table names, so a typo can create an unused table instead of failing.

Minimal Local Contract Example

Put the following example in an importable custom_model_demo.py in a scratch directory, then run python custom_model_demo.py with this checkout installed. It intentionally returns input text, not speech recognition. It does not need weights, audio, a hub download, or GPU. The parameter is present because this AutoModel.inference inspects next(model.parameters()).device; a parameterless toy fails that path.

import torch
from funasr import AutoModel
from funasr.register import tables

MODEL_NAME = "DocsEchoModelV1"
if MODEL_NAME in tables.model_classes:
    raise RuntimeError(f"Registry collision: {MODEL_NAME}")


@tables.register("model_classes", MODEL_NAME)
class DocsEchoModel(torch.nn.Module):
    def __init__(self, **kwargs):
        super().__init__()
        self.anchor = torch.nn.Parameter(torch.zeros(1), requires_grad=False)

    def inference(
        self, data_in, data_lengths=None, key=None,
        tokenizer=None, frontend=None, **kwargs,
    ):
        results = [
            {"key": sample_key, "text": str(value)}
            for sample_key, value in zip(key, data_in)
        ]
        return results, {}


if __name__ == "__main__":
    model = AutoModel(
        model=MODEL_NAME,
        model_conf={},
        device="cpu",
        disable_update=True,
        disable_pbar=True,
    )
    result = model.generate(input=["hello", "world"], data_type="text")
    assert [row["text"] for row in result] == ["hello", "world"]
    assert all(isinstance(row["key"], str) for row in result)
    print([row["text"] for row in result])

Expected final printed list: ['hello', 'world']. This is a wiring check only. The focused documentation test executes this block in a temporary file against the checked-out implementation; it does not validate a speech model.

The presence of model_conf={} is intentional: AutoModel.build_model skips hub/config resolution when that key is supplied. Thus model is the already-registered class key here, not a model directory or hub ID. Import the custom module yourself first; adding remote_code to this direct path does not cause it to be imported. Resolved kwargs are merged over model_conf and passed to the constructor, including built tokenizer/frontend objects, device and vocabulary/input sizes. Accept appropriate named arguments plus **kwargs, as the real models do.

Inference and Training Contracts

Interface Contract in this checkout
Model object Normally torch.nn.Module; must support .to(...), .eval() and .parameters(). Constructor configuration is model-specific.
inference inputs Without VAD, AutoModel batches input into data_in and key lists and calls model.inference(**batch, **kwargs) under torch.no_grad(). With data_type="fbank" and one item, it passes the feature object directly and supplies data_lengths=input_len. Tokenizer/frontend are resolved objects or None.
Model-level return Return a two-item tuple (results, meta_data): results is list[dict], meta_data is a dictionary. For ASR use string key and text fields, preserving the input order/keys. Returning a bare list of result dictionaries is wrong: AutoModel treats its first item as the entire batch result.
Metadata Optional load_data, extract_feat, batch_data_time. For audio, batch_data_time is a positive duration in seconds, not milliseconds; zero causes division by zero in timing code. Omitting it gives the internal -1 sentinel, as in the non-audio toy, not a meaningful speed measurement.
Public return AutoModel.generate(...) returns the flattened result list. Additional fields such as timestamps are model-specific; registration does not promise them. VAD, punctuation, speaker and streaming integrations need additional compatible behavior and independent tests.
Training Implement differentiable forward matching the dataset collator's named tensor fields. Trainer unpacks (loss, stats, weight); SenseVoice and Nano use force_gatherable. The toy deliberately has no training forward.
Export Export utility calls the model's export and then model-specific methods such as export_dummy_inputs, input/output names and dynamic axes. Registration alone does not implement these.

A custom model must itself implement audio loading, feature preparation, tokenization/decoding and batching where needed. Copy the contract from a close real model, not its capabilities without implementation. For training, also implement/configure the appropriate dataset and loss; see training.

Loading Reviewed Custom Code and Weights

There are two distinct routes:

  1. Direct registration: import your module; pass its key and model_conf as in the toy. Provide compatible tokenizer/frontend/configuration and an existing init_param if weights are needed. No hub code import is performed on this path.
  2. Model-directory resolution: pass a reviewed local directory or hub ID without model_conf. The loader reads configuration.json file metadata or config.yaml, resolves the model key/assets and loads weights. A simple local config.yaml directory normally also needs model.pt plus all referenced tokenizer/frontend assets. Arbitrary HF weight folders are not automatically FunASR model directories.

For the second route, this is the ModelScope-path interface, not a self-contained runnable example. models/custom-asr must already contain the compatible reviewed configuration and weights, and custom_asr_model.py must register the exact configuration key:

from funasr import AutoModel

model = AutoModel(
    model="./models/custom-asr",
    hub="ms",
    trust_remote_code=True,
    remote_code="./custom_asr_model.py",
    device="cpu",
    disable_update=True,
)
print(model.generate(input="data/audio/heldout.wav"))

The current download_from_ms imports remote_code when trust_remote_code=True, defaulting to module name model if omitted. The importer supports module/file paths (and URL download), appends the directory to sys.path and imports by basename. A relative path is relative to the working directory, not automatically to the weights directory. Basename collisions and Python's import cache can select an already-loaded module; use distinct module names and verify the active class. Import exceptions are printed rather than re-raised by this helper, so inspect errors and registry state.

The hub="hf" path differs: in this checkout download_from_hf can install model requirements under the trust flag but does not call import_module_from_path. Do not assume remote_code is executed there. Import your reviewed custom module explicitly before constructing with hub="hf", or use the direct-registration route with complete configuration. The local config.yaml fallback also handles init_param differently between hubs: ModelScope retains an existing explicit path; Hugging Face assigns the directory's model.pt. Check the resolved path rather than assuming a selected checkpoint override survived.

Real Examples and Boundaries

  • SenseVoiceSmall implementation demonstrates a registered model with training, inference and export integration. Retain its model-specific configuration and tokenizer assumptions.
  • Nano demo1.py shows trust_remote_code=True, remote_code="./model.py", hub="ms"; it assumes the recipe directory. Its local implementation registers FunASRNano and imports sibling ctc and tools modules. That can overwrite the built-in implementation; the two are not interchangeable for every feature, including built-in LoRA. Audit the active class and checkpoint keys.
  • MOSS adapter integrates the third-party OpenMOSS model, and explicitly rejects training in forward. It is not evidence that a registered model supports fine-tuning or export.
  • The original registration tutorial and general tutorial remain useful historical references; use the current source behavior above when their examples differ.

Safety and Verification

trust_remote_code=True permits Python execution; the hub loaders can also install a model directory's requirements.txt. A local directory is not inherently trustworthy. Review source, dependencies and weight serialization, use isolated environments, and do not load untrusted pickled checkpoints. Never interpolate untrusted URLs, module names or configs into this workflow. Keep a reviewed local snapshot with hashes when remote revision handling cannot be relied on; model_revision is not a universal pin across all loader paths here.

Check the active registry key/class/source before loading. Confirm exact model/config/tokenizer compatibility and inspect missing/unexpected weights: ignore_init_mismatch defaults to true in AutoModel, and a nonexistent direct init_param prints an error rather than guaranteeing construction fails. Validate checkpoint existence yourself. Test one item, multiple items, error handling and the intended pipeline before training/export/deployment. FunASR software's MIT license does not replace model or upstream component licenses.

Focused syntax, repository-link and no-download toy-contract checks:

python -m pytest -q tests/test_training_docs_contract.py

These checks do not certify arbitrary custom code, actual ASR quality, GPU training, real-checkpoint restoration or export compatibility.