Loading¶

trimbed.loading ¶

Loading tokenizers and models.

load_tokenizer ¶

load_tokenizer(
    model: str,
    revision: str | None = None,
    trust_remote_code: bool = False,
) -> PreTrainedTokenizerFast

Load a tokenizer as a fast tokenizer, converting it if necessary.

Tokenizers distributed only as a SentencePiece model (mT5's spiece.model, say) are converted to the tokenizers format on load so we can trim them following the same code path as everything else.

Parameters:

Name	Type	Description	Default
`model`	`str`	Hub model id or local path, e.g. `"codefuse-ai/F2LLM-v2-160M"` or `"./trimmed/f2llm-nl"`.	required
`revision`	`str \| None`	Optional revision to pin, e.g. `"refs/pr/1"` or a commit sha.	`None`
`trust_remote_code`	`bool`	Allow tokenizer code shipped with the checkpoint.	`False`

Returns:

Type	Description
`PreTrainedTokenizerFast`	A `transformers` fast tokenizer, e.g. a `Qwen2Tokenizer` for codefuse-ai/F2LLM-v2-160M or a `BertTokenizer` for google-bert/bert-base-cased.

Raises:

Type	Description
`MissingDependencyError`	If conversion needs `sentencepiece`/`protobuf`.
`ValueError`	If the result is not backed by a `tokenizers.Tokenizer`.

View source on GitHub: src/trimbed/loading.py lines 23–60

require_torch ¶

require_torch() -> ModuleType

Import and return torch or explain how to install it.

Returns:

Type	Description
`ModuleType`	The imported `torch` module.

Raises:

Type	Description
`MissingDependencyError`	If torch is not installed.

View source on GitHub: src/trimbed/loading.py lines 63–76

resolve_model_class ¶

resolve_model_class(
    model: str,
    revision: str | None = None,
    config: EmbeddingTrimConfig | None = None,
    trust_remote_code: bool = False,
) -> type[PreTrainedModel] | type[AutoModel]

Find the model class to load a checkpoint with.

The class named in config.architectures is what the checkpoint actually contains, so that is what gets loaded.

Parameters:

Name	Type	Description	Default
`model`	`str`	Hub model id or local path.	required
`revision`	`str \| None`	Optional revision to pin.	`None`
`config`	`EmbeddingTrimConfig \| None`	Embedding-trimming settings. Set `auto_class` to override the inference.	`None`
`trust_remote_code`	`bool`	Allow modelling code shipped with the checkpoint.	`False`

Returns:

Type Description

A class exposing `from_pretrained`, read off `config.architectures`

e.g. Qwen3ForCausalLM for Qwen/Qwen3-0.6B, BertForMaskedLM for google-bert/bert-base-cased, T5ForConditionalGeneration for google-t5/t5-small, and Qwen3Model for codefuse-ai/F2LLM-v2-160M, which really is a base model with no head (it is an embedding model that uses last-token pooling). The second arm of the return type is the AutoModel fallback for a remote-code checkpoint whose class transformers does not export, which is separate because the auto classes are factories rather than PreTrainedModel subclasses.

Raises:

Type	Description
`ValueError`	If `EmbeddingTrimConfig.auto_class` names something transformers does not export, e.g. a typo like `"AutoModelForCasualLM"`, although perhaps it would not be bad if we'd have some more casual lm models. :-)

View source on GitHub: src/trimbed/loading.py lines 79–131

load_model ¶

load_model(
    model: str,
    revision: str | None = None,
    config: EmbeddingTrimConfig | None = None,
    trust_remote_code: bool = False,
) -> PreTrainedModel

Load a model for embedding trimming.

Parameters:

Name	Type	Description	Default
`model`	`str`	Hub model id or local path.	required
`revision`	`str \| None`	Optional revision to pin.	`None`
`config`	`EmbeddingTrimConfig \| None`	Embedding-trimming settings, which control the class, dtype and placement, e.g. `dtype="bfloat16"` with `device="cuda"`.	`None`
`trust_remote_code`	`bool`	Allow modelling code shipped with the checkpoint, as the gte and jina encoders need.	`False`

Returns:

Type	Description
`PreTrainedModel`	A `transformers` model in eval mode.

View source on GitHub: src/trimbed/loading.py lines 134–164