Read OSS

Pipeline、Tokenizer、そして Transformers の拡張

中級

前提知識

  • 第1回: インポートシステム
  • 第2回: Auto クラス
  • トークナイズの基本的な理解(BPE、SentencePiece)
  • NLP・ビジョン・音声の推論タスクへの馴染み

Pipeline、Tokenizer、そして Transformers の拡張

このシリーズでは、Transformers の遅延インポートシステムから始まり、モデル解決、重みのロード、生成、学習まで順に追ってきました。最終回となるこの記事では、モデルの上下に位置するレイヤーを取り上げます。具体的には、推論をワンライナーで実現する高レベルの pipeline() API、3つのバックエンドにわたってテキスト前処理を担うトークナイザの階層、マルチモーダルモデル向けの ProcessorMixin、新しいモデルをライブラリへコントリビュートするための拡張ポイントです。

pipeline() ファクトリ

pipeline() は Transformers の中で最もシンプルな API です。タスク文字列を適切な Pipeline サブクラスに解決し、モデルとトークナイザを自動選択したうえで、呼び出し可能なオブジェクトを返します。

generator = pipeline("text-generation", model="meta-llama/Llama-2-7b-hf")
output = generator("Once upon a time")
flowchart TD
    A["pipeline('text-generation',<br/>model='...')"] --> B["Resolve task → Pipeline class"]
    B --> C["PipelineRegistry lookup"]
    C --> D["TextGenerationPipeline"]
    D --> E["Auto-select model<br/>(AutoModelForCausalLM)"]
    E --> F["Auto-select tokenizer<br/>(AutoTokenizer)"]
    F --> G["Load model<br/>from_pretrained()"]
    G --> H["Return configured<br/>Pipeline instance"]

関数のシグネチャは taskmodeltokenizerconfigprocessordevicedevice_mapdtypetrust_remote_code を受け付けます。task だけを指定するとライブラリがデフォルトのモデルを選び、model を指定するとモデルの config からタスクを推論します。

解決処理は PipelineRegistry が担っており、タスク文字列("text-generation""text-classification""image-to-text" など)を Pipeline サブクラスとデフォルトモデルにマッピングします。エイリアスもサポートされており、たとえば "sentiment-analysis"TextClassificationPipeline に対応しています。

Pipeline 基底クラス: preprocess → _forward → postprocess

すべての pipeline は Pipeline を継承しており、3ステップのパターンを定義しています。

sequenceDiagram
    participant User
    participant Pipe as Pipeline.__call__
    participant Pre as preprocess()
    participant Fwd as _forward()
    participant Post as postprocess()

    User->>Pipe: pipe("Hello world")
    Pipe->>Pre: preprocess("Hello world")
    Pre-->>Pipe: model_inputs (dict of tensors)
    Pipe->>Fwd: _forward(model_inputs)
    Note over Fwd: model(**inputs)<br/>with torch.no_grad()
    Fwd-->>Pipe: model_outputs
    Pipe->>Post: postprocess(model_outputs)
    Post-->>User: [{"generated_text": "..."}]

基底クラスの Pipeline.__call__ はバッチ処理、長い入力のチャンク分割、no-grad コンテキストの管理を担います。サブクラスはこの3つのメソッドをオーバーライドします。

  1. preprocess() — 生の入力(文字列、画像、音声)をモデルが受け付けるテンソルに変換する
  2. _forward() — モデルを実行する(デフォルト: self.model(**inputs)
  3. postprocess() — モデルの出力を人間が読める形式に変換する

どの前処理コンポーネントをロードするかはクラスのフラグで制御されます。

class TextGenerationPipeline(Pipeline):
    _load_tokenizer = True
    _load_image_processor = False
    _load_feature_extractor = False
    _pipeline_calls_generate = True  # Uses model.generate() instead of model()

特に重要なのが _pipeline_calls_generate = True フラグです。テキスト生成 pipeline はモデルの生の forward パスではなく model.generate()(第5回)を呼び出すため、デコード、キャッシング、ストリーミングの仕組みをそのまま利用できます。

ヒント: バッチ処理を伴う本番推論では pipe(texts, batch_size=16) を使いましょう。パディング、照合、バッチの反復処理は Pipeline が内部で行います。ストリーミングが必要な場合は、TextStreamer を pipeline の generation kwargs として渡してください。

トークナイザの階層

Transformers のトークナイズは、複数のアーキテクチャ変遷を経てきました。現在のシステムは共通インターフェースのもとで3つのバックエンド実装を提供しています。

classDiagram
    class PreTrainedTokenizerBase {
        +__call__(text, ...)
        +encode(text)
        +decode(ids)
        +from_pretrained()
        +save_pretrained()
        +padding_side: str
        +model_max_length: int
    }
    class PreTrainedTokenizerPython {
        «Pure Python BPE/WordPiece»
        +_tokenize(text) → list[str]
        +_convert_token_to_id(token) → int
    }
    class PreTrainedTokenizerTokenizers {
        «Rust tokenizers backend»
        +_tokenizer: tokenizers.Tokenizer
        +_batch_encode_plus()
    }
    class PreTrainedTokenizerSentencePiece {
        «SentencePiece backend»
        +sp_model: SentencePieceProcessor
    }
    PreTrainedTokenizerBase <|-- PreTrainedTokenizerPython
    PreTrainedTokenizerBase <|-- PreTrainedTokenizerTokenizers
    PreTrainedTokenizerBase <|-- PreTrainedTokenizerSentencePiece

PreTrainedTokenizerBase は、パディング、トランケーション、アテンションマスクの生成、返り値の型選択(pt / np / tf)を扱う共通の __call__ インターフェースを定義しています。バッチエンコードの組み合わせロジックすべてを網羅した約2000行のファイルです。

3つのバックエンドはそれぞれ次のとおりです。

  • tokenization_python.py — BPE、WordPiece、Unigram の Pure Python 実装。最も低速ですが、依存ライブラリが不要です。
  • tokenization_utils_tokenizers.pytokenizers Rust ライブラリをラップし、10〜100倍の高速化を実現します。これがデフォルトの「fast」トークナイザです。
  • tokenization_utils_sentencepiece.py — Google の SentencePiece ライブラリをラップします。LLaMA、T5 をはじめ多くのモデルで採用されています。

第1回で触れたように、モジュールエイリアスシステムは後方互換性のために tokenization_utils_fasttokenization_utils_tokenizers へのマッピングを行っています。AutoTokenizer クラスは利用可能な最速のバックエンドを選択し、tokenizers ライブラリがインストールされている場合は Rust トークナイザを優先します。

ヒント: モデル固有のトークナイザクラスを直接使うのではなく、常に AutoTokenizer.from_pretrained(model_name) を使いましょう。Auto クラスがバックエンドの選択を行い、モデルの語彙との互換性を確保してくれます。

マルチモーダルモデル向け ProcessorMixin

LLaVA、Whisper、CLIP のような現代のモデルは、テキスト以外のモダリティも扱います。ProcessorMixin は、トークナイザとモダリティ固有のプロセッサを組み合わせた統一 __call__ インターフェースを提供します。

classDiagram
    class ProcessorMixin {
        +tokenizer: PreTrainedTokenizerBase
        +image_processor: BaseImageProcessor
        +feature_extractor: FeatureExtractionMixin
        +__call__(*args, **kwargs)
        +from_pretrained()
        +save_pretrained()
    }
    class LlavaProcessor {
        +tokenizer: LlamaTokenizer
        +image_processor: CLIPImageProcessor
        +__call__(text, images)
    }
    class WhisperProcessor {
        +tokenizer: WhisperTokenizer
        +feature_extractor: WhisperFeatureExtractor
        +__call__(audio, text)
    }
    ProcessorMixin <|-- LlavaProcessor
    ProcessorMixin <|-- WhisperProcessor

processor の __call__ メソッドは一般的に次の順で動作します。

  1. 対応するモダリティプロセッサで画像や音声を処理する
  2. テキストをトークナイズする
  3. 正しい shape で単一の BatchFeature dict にまとめる

ProcessorMixin はシリアライズも担当しています。save_pretrained() でトークナイザとプロセッサの config を両方保存し、from_pretrained() でそれらを読み込みます。

拡張ポイント: 新しいモデルの追加

このシリーズで解説してきた内容を踏まえ、Transformers に新しいモデルをコントリビュートする際の完全なチェックリストを以下に示します。

flowchart TD
    A["1. Create config<br/>with model_type"] --> B["2. Create model<br/>inheriting PreTrainedModel"]
    B --> C["3. Create tokenizer<br/>(or reuse existing)"]
    C --> D["4. Register in Auto mappings<br/>CONFIG_MAPPING_NAMES<br/>MODEL_FOR_*_MAPPING_NAMES"]
    D --> E["5. Add __all__ exports<br/>to each .py file"]
    E --> F["6. Add TYPE_CHECKING block<br/>in model __init__.py"]
    F --> G["7. Write tests<br/>and documentation"]

ステップ 1: Configuration。 configuration_<model>.py を作成し、model_type を設定してすべてのハイパーパラメータをデフォルト値付きの型付きフィールドとして宣言するクラスを定義します。バリデーションには @strict を使いましょう。モデルが並列化をサポートする場合は base_model_tp_planbase_model_pp_plan も宣言してください。

ステップ 2: Model。 modeling_<model>.py を作成します。PreTrainedModel を継承し、クラスフラグ(_supports_sdpa など)を設定して、アテンションのディスパッチには ALL_ATTENTION_FUNCTIONS.get_interface() を使った forward パスを実装します。デコーダ層には GradientCheckpointingLayer を使いましょう。タスクヘッドは多重継承で汎用クラス(GenericForSequenceClassification など)を利用してください。

ステップ 3: Auto 登録。 models/auto/ 内の CONFIG_MAPPING_NAMES と該当する MODEL_FOR_*_MAPPING_NAMES dict にエントリを追加します。手動登録が必要なのはここだけで、遅延インポートシステム(第1回)は __all__ エクスポートからその他すべてを自動的に検出します。

ステップ 4: Tokenizer。 新しいトークナイザを作成するか、既存のものを再利用します。最近の LLM の多くは LlamaTokenizer や SentencePiece ベースのトークナイザを再利用しています。

trust_remote_code: 動的モジュールのロード

Transformers にまだマージされていないモデルに対しては、dynamic_module_utils.py が Hub リポジトリから任意の Python コードをロードする仕組みを提供しています。

model = AutoModelForCausalLM.from_pretrained(
    "org/new-model",
    trust_remote_code=True
)

これは Hub からモデルの Python ファイルをダウンロードし、サンドボックス化されたモジュール名前空間で実行したうえで、Auto システムにクラスを登録します。config.jsonauto_map フィールドが、各 Auto クラスに対してどのリモートクラスを使うかを Transformers に伝えます。

{
    "model_type": "new_model",
    "auto_map": {
        "AutoConfig": "configuration_new_model.NewModelConfig",
        "AutoModelForCausalLM": "modeling_new_model.NewModelForCausalLM"
    }
}

CLI スキャフォールディングツール

cli/add_new_model_like.py スクリプトは、既存のモデルをコピー・改変することで新しいモデルのボイラープレートを生成します。ディレクトリ構造の作成、クラスのリネーム、Auto マッピングの更新、テストファイルの生成といった面倒な作業を自動化してくれます。

Config に宣言する並列化プラン

第3回で見たように、LLaMA はテンソル並列プランを config 内で直接宣言しています。

base_model_tp_plan = {
    "layers.*.self_attn.q_proj": "colwise",
    "layers.*.self_attn.k_proj": "colwise",
    "layers.*.self_attn.v_proj": "colwise",
    "layers.*.self_attn.o_proj": "rowwise",
    "layers.*.mlp.gate_proj": "colwise",
    "layers.*.mlp.up_proj": "colwise",
    "layers.*.mlp.down_proj": "rowwise",
}

これにより、テンソル並列化のためにコードを一切変更する必要がなくなります。model.tensor_parallel(device_mesh) を呼び出すだけで、ライブラリがこのプランに従って各層に列方向・行方向のシャーディングを適用します。同様に、base_model_pp_plan は入出力テンソルの明示的な仕様とともに、どの層をどのパイプライン並列ステージに配置するかを宣言します。

全体アーキテクチャ

シリーズの締めくくりとして、すべてのピースがどのように組み合わさるかを示します。

flowchart TD
    subgraph "User API"
        A["pipeline()"] 
        B["AutoModelForCausalLM"]
        C["Trainer"]
    end

    subgraph "Resolution Layer"
        D["_LazyModule"]
        E["Auto Mappings"]
        F["CONFIG_MAPPING"]
    end

    subgraph "Model Layer"
        G["PreTrainedModel"]
        H["AttentionInterface"]
        I["Generic Head Classes"]
    end

    subgraph "Weight Layer"
        J["from_pretrained()"]
        K["WeightConverter"]
        L["HfQuantizer"]
    end

    subgraph "Generation"
        M["GenerationMixin"]
        N["Cache System"]
        O["LogitsProcessors"]
    end

    subgraph "Training"
        P["Trainer Loop"]
        Q["Callbacks"]
        R["Distributed Backends"]
    end

    subgraph "Preprocessing"
        S["Tokenizers"]
        T["ProcessorMixin"]
    end

    A --> E
    B --> E
    E --> D
    D --> G
    G --> H
    G --> I
    J --> K
    J --> L
    G --> M
    M --> N
    M --> O
    C --> P
    P --> Q
    P --> R
    A --> S
    A --> T

この図のすべてのボックスは、このシリーズを通じて詳しく見てきたコードに対応しています。インポートシステム(第1回)が土台を提供し、Auto クラス(第2回)がディスパッチを担い、モデルの構造(第3回)が計算を実行し、重みのロード(第4回)が初期化を行い、生成(第5回)が推論を担い、Trainer(第6回)が学習を支えます。そして今回の記事が、それらをひとつに結びつけるユーザー向け API と拡張ポイントをカバーしています。

まとめ

Transformers は、一貫したパターンによって卓越した複雑さを管理しているライブラリです。遅延ロードの徹底、宣言的な設定、composable な pipeline、そして高レベル API から低レベル実装への段階的な委譲がその柱です。コードベースは読む価値があります。ほぼすべての設計上の決定には明確な根拠があり、450以上のアーキテクチャ、十数種のバックエンド、数百万のユーザーをサポートするという制約から生まれています。

このシリーズを通じて追ってきたパターンは、あらゆる大規模ライブラリ設計に応用できます。遅延インポートのための _LazyModule、遅延クラス解決のための _LazyAutoMapping、プラガブルなバックエンドのための AttentionInterface、2段階ディスパッチのための GeneralInterface などがその例です。汎用ヘッドクラスによるコード再利用や、WeightConverter によるチェックポイント互換性も重要なパターンです。これらを理解することは、Transformers をより上手く使えるようになるだけでなく、より優れたシステムプログラマになることにもつながります。