Create an input processor for a specific model. Args: model_path_or_dir: Path or repo id used to locate pretrained config/tokenizer. tokenizer: Tokenizer instance. checkpoint_format: Checkpoint format identifier. "HF" uses Hugging Face-style config loading; a
(
model_path_or_dir: str,
tokenizer,
checkpoint_format: Optional[str] = "HF",
)
| 588 | |
| 589 | |
| 590 | def create_input_processor( |
| 591 | model_path_or_dir: str, |
| 592 | tokenizer, |
| 593 | checkpoint_format: Optional[str] = "HF", |
| 594 | ) -> Union[InputProcessor, BaseMultimodalInputProcessor]: |
| 595 | """Create an input processor for a specific model. |
| 596 | |
| 597 | Args: |
| 598 | model_path_or_dir: Path or repo id used to locate pretrained config/tokenizer. |
| 599 | tokenizer: Tokenizer instance. |
| 600 | checkpoint_format: Checkpoint format identifier. "HF" uses Hugging Face-style |
| 601 | config loading; any other value skips HF config loading. Default is "HF". |
| 602 | |
| 603 | Returns: |
| 604 | An InputProcessor implementation (model-specific if registered; otherwise DefaultInputProcessor). |
| 605 | """ |
| 606 | from tensorrt_llm._torch.model_config import ModelConfig |
| 607 | from tensorrt_llm._torch.models import get_model_architecture |
| 608 | |
| 609 | config = None |
| 610 | |
| 611 | if checkpoint_format == "HF": |
| 612 | try: |
| 613 | model_config = ModelConfig.from_pretrained(model_path_or_dir, |
| 614 | trust_remote_code=True) |
| 615 | config = model_config.pretrained_config |
| 616 | except (ValueError, EnvironmentError) as e: |
| 617 | logger.debug( |
| 618 | f"Unable to load HF config from {model_path_or_dir}: {e}. Falling back." |
| 619 | ) |
| 620 | elif checkpoint_format in ("mistral", "mistral_large_3"): |
| 621 | logger.debug(f"Detected checkpoint_format={checkpoint_format}.") |
| 622 | from tensorrt_llm._torch.models.checkpoints.mistral.config_loader import \ |
| 623 | MistralConfigLoader |
| 624 | model_config = MistralConfigLoader().load(model_path_or_dir) |
| 625 | config = model_config.pretrained_config |
| 626 | else: |
| 627 | logger.debug( |
| 628 | f"checkpoint_format={checkpoint_format}; skipping HF config load.") |
| 629 | |
| 630 | if config is not None: |
| 631 | try: |
| 632 | model_cls, _ = get_model_architecture(config) |
| 633 | input_processor_cls = INPUT_PROCESSOR_REGISTRY._input_processors_cls_by_model_type \ |
| 634 | .get(model_cls) |
| 635 | except RuntimeError: # unregistered model |
| 636 | logger.info("Unregistered model, using DefaultInputProcessor") |
| 637 | input_processor_cls = None |
| 638 | if input_processor_cls is not None: |
| 639 | return input_processor_cls(model_path_or_dir, |
| 640 | config, |
| 641 | tokenizer, |
| 642 | trust_remote_code=True) |
| 643 | |
| 644 | return DefaultInputProcessor(None, None, tokenizer) |
| 645 | |
| 646 | |
| 647 | def create_input_processor_with_hash( |