自然言語処理（NLP）におけるHugging Faceと分散処理基盤Sparkの統合手法について解説する。Spark 3.4以降の新機能活用と、Spark 3.3.1向けカスタム実装の2種類のアプローチを紹介する。

Spark 3.4以降でのネイティブ統合

Sparkのバージョン3.4以上では組み込みの機械学習サポートが強化されており、分散環境でのモデル推論が容易に実装できる。

実装の要点

• ワーカー単位のモデルロード戦略
• モデルディレクトリのクリーンアップ事前処理

Spark 3.3.1向けカスタム実装

既存環境のアップグレードが困難な場合、分散推論機能を独自に実装する必要がある。

推論モデルキャッシュ機構

from threading import Lock
from uuid import UUID
from collections import OrderedDict

class InferenceModelCache:
    _storage = OrderedDict()
    _capacity = 3
    _lock = Lock()

    @staticmethod
    def register(model_id: UUID, infer_func: callable):
        with InferenceModelCache._lock:
            InferenceModelCache._storage[model_id] = infer_func
            InferenceModelCache._storage.move_to_end(model_id)
            if len(InferenceModelCache._storage) > InferenceModelCache._capacity:
                InferenceModelCache._storage.popitem(last=False)

    @staticmethod
    def fetch(model_id: UUID) -> Optional[callable]:
        with InferenceModelCache._lock:
            func = InferenceModelCache._storage.get(model_id)
            if func:
                InferenceModelCache._storage.move_to_end(model_id)
            return func

分散推論処理実装

import uuid
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import ArrayType, StringType
import pandas as pd
import numpy as np

def create_inference_udf(model_initializer, output_schema, batch_size=100):
    model_id = uuid.uuid4()
    
    def batch_inference(data_iterator):
        from inference_model_cache import InferenceModelCache
        inference_func = InferenceModelCache.fetch(model_id)
        if not inference_func:
            inference_func = model_initializer()
            InferenceModelCache.register(model_id, inference_func)
        
        for data_batch in data_iterator:
            for mini_batch in split_into_batches(data_batch, batch_size):
                predictions = inference_func(mini_batch)
                yield transform_predictions(predictions, output_schema)
    
    return pandas_udf(batch_inference, output_schema)

def split_into_batches(data: pd.DataFrame, size: int):
    idx = 0
    while idx < len(data):
        yield data.iloc[idx:idx+size]
        idx += size

def transform_predictions(results, schema):
    # 予測結果のスキーマ変換ロジック
    return pd.DataFrame(results)

Hugging Faceモデル統合例

from transformers import AutoTokenizer, AutoModel

def setup_embedding_model():
    import os
    os.system("rm -rf ./bert-model")
    os.system("hadoop fs -get /models/bert-base-chinese")
    
    tokenizer = AutoTokenizer.from_pretrained("./bert-model")
    model = AutoModel.from_pretrained("./bert-model")
    
    def generate_embeddings(texts):
        inputs = tokenizer(texts.tolist(), padding=True, truncation=True, max_length=128, return_tensors="pt")
        outputs = model(**inputs)
        embeddings = outputs.last_hidden_state.mean(dim=1).detach().numpy()
        return [embedding.astype(np.float32).tolist() for embedding in embeddings]
    
    return generate_embeddings

# UDF登録
embedding_udf = create_inference_udf(
    setup_embedding_model,
    ArrayType(ArrayType(FloatType())),
    batch_size=64
)

# 推論実行
df.withColumn("embeddings", embedding_udf("text_column"))

手法比較