在自然语言处理（NLP）领域，Hugging Face 是不可或缺的处理库，而 Spark 则是大数据处理的必备工具。将两者的优势结合起来，可以实现高效的 NLP 大数据处理。以下是结合 Hugging Face 和 Spark 的两种方法，基于 Spark & PySpark 3.3.1 版本进行探索。

方法一：升级 Spark 版本至 3.4 及以上

如果你愿意升级 Spark 版本到 3.4 或更高版本，那么结合 Hugging Face 和 Spark 将变得非常方便。Spark 3.4 及以上版本天然支持加载模型进行预测。

关键步骤说明：

1. 模型加载策略：需要为每个 Worker 单独加载模型，确保模型在分布式环境中的可用性。
2. 文件夹管理：在加载 Hugging Face 预训练模型之前，务必删除之前的模型文件夹，防止加载失败。

注：如果图片无法显示，请检查链接合法性或稍后重试。

方法二：基于 Spark 3.3.1 的手动封装接口

如果你希望保持当前的 Spark 3.3.1 版本，那么可以通过手动封装接口来实现 Hugging Face 和 Spark 的结合。以下是详细的代码实现和关键说明。

封装分布式的模型缓存

为了高效管理模型加载和缓存，我们从spark3.4的源代码中抽取了一个分布式的模型缓存机制：

from collections import OrderedDict
from threading import Lock
from typing import Callable, Optional
from uuid import UUID


class ModelCache:
    """Cache for model prediction functions on executors.

    This requires the `spark.python.worker.reuse` configuration to be set to `true`, otherwise a
    new python worker (with an empty cache) will be started for every task.

    If a python worker is idle for more than one minute (per the IDLE_WORKER_TIMEOUT_NS setting in
    PythonWorkerFactory.scala), it will be killed, effectively clearing the cache until a new python
    worker is started.

    Caching large models can lead to out-of-memory conditions, which may require adjusting spark
    memory configurations, e.g. `spark.executor.memoryOverhead`.
    """

    _models: OrderedDict = OrderedDict()
    _capacity: int = 3  # "reasonable" default size for now, make configurable later, if needed
    _lock: Lock = Lock()

    @staticmethod
    def add(uuid: UUID, predict_fn: Callable) -> None:
        with ModelCache._lock:
            ModelCache._models[uuid] = predict_fn
            ModelCache._models.move_to_end(uuid)
            if len(ModelCache._models) > ModelCache._capacity:
                ModelCache._models.popitem(last=False)

    @staticmethod
    def get(uuid: UUID) -> Optional[Callable]:
        with ModelCache._lock:
            predict_fn = ModelCache._models.get(uuid)
            if predict_fn:
                ModelCache._models.move_to_end(uuid)
            return predict_fn

封装处理逻辑

from __future__ import annotations
import os
import argparse
import random
import logging
import pandas as pd
from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, column, encode
from pyspark.sql.types import *
from datetime import datetime, timedelta
import requests as req
from io import BytesIO
import numpy as np
import uuid
import inspect
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import (
    ArrayType,
    ByteType,
    DataType,
    DoubleType,
    FloatType,
    IntegerType,
    LongType,
    ShortType,
    StringType,
    StructType,
)
from typing import Any, Callable, Iterator, List, Mapping, TYPE_CHECKING, Tuple, Union, Optional

supported_scalar_types = (
    ByteType,
    ShortType,
    IntegerType,
    LongType,
    FloatType,
    DoubleType,
    StringType,
)

hadoop = os.path.join(os.environ['HADOOP_COMMON_HOME'], 'bin/hadoop')

def init_spark():
    """初始化 SparkSession 配置"""
    spark = SparkSession.builder \
        .config("spark.sql.caseSensitive", "false") \
        .config("spark.shuffle.spill", "true") \
        .config("spark.shuffle.spill.compress", "true") \
        .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer") \
        .config("metastore.catalog.default", "hive") \
        .config("spark.sql.hive.convertMetastoreOrc", "true") \
        .config("spark.kryoserializer.buffer.max", "1024m") \
        .config("spark.kryoserializer.buffer", "64m") \
        .config("spark.driver.maxResultSize","4g") \
        .config("spark.sql.broadcastTimeout", "36000") \
        .enableHiveSupport() \
        .getOrCreate()
    return spark

def system_command(command):
    """执行系统命令"""
    code = os.system(command)
    if code != 0:
        logging.error(f"Command: ({command}) execute failed.")
    else:
        logging.info(f"Command: ({command}) execute succeed.")

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(usage="it's usage tip.",
                                     description="user tags prefer")
    parser.add_argument("--db", default="", help="hive表")
    parser.add_argument("--date", default="", help="日期")
    parser.add_argument("--output_path", default="", help="输出路径")
    parser.add_argument("--batch_size", default=16, help="输出路径")
    return parser.parse_args()

def _batched(
    data: Union[pd.Series, pd.DataFrame, Tuple[pd.Series]], batch_size: int
) -> Iterator[pd.DataFrame]:
    """将 pandas dataframe/series 分批处理"""
    if isinstance(data, pd.DataFrame):
        df = data
    elif isinstance(data, pd.Series):
        df = pd.concat((data,), axis=1)
    else:  # isinstance(data, Tuple[pd.Series])
        df = pd.concat(data, axis=1)

    index = 0
    data_size = len(df)
    while index < data_size:
        yield df.iloc[index : index + batch_size]
        index += batch_size

def _is_tensor_col(data: Union[pd.Series, pd.DataFrame]) -> bool:
    """检查数据是否为张量列"""
    if isinstance(data, pd.Series):
        return data.dtype == np.object_ and isinstance(data.iloc[0], (np.ndarray, list))
    elif isinstance(data, pd.DataFrame):
        return any(data.dtypes == np.object_) and any(
            [isinstance(d, (np.ndarray, list)) for d in data.iloc[0]]
        )
    else:
        raise ValueError(
            "Unexpected data type: {}, expected pd.Series or pd.DataFrame.".format(type(data))
        )

def _has_tensor_cols(data: Union[pd.Series, pd.DataFrame, Tuple[pd.Series]]) -> bool:
    """检查输入是否包含张量值列"""
    if isinstance(data, (pd.Series, pd.DataFrame)):
        return _is_tensor_col(data)
    else:  # isinstance(data, Tuple)
        return any(_is_tensor_col(elem) for elem in data)

def _validate_and_transform_multiple_inputs(
    batch: pd.DataFrame, input_shapes: List[Optional[List[int]]], num_input_cols: int
) -> List[np.ndarray]:
    """验证并转换多个输入"""
    multi_inputs = [batch[col].to_numpy() for col in batch.columns]
    if input_shapes:
        if len(input_shapes) == num_input_cols:
            multi_inputs = [
                np.vstack(v).reshape([-1] + input_shapes[i])  # type: ignore
                if input_shapes[i]
                else v
                for i, v in enumerate(multi_inputs)
            ]
            if not all([len(x) == len(batch) for x in multi_inputs]):
                raise ValueError("Input data does not match expected shape.")
        else:
            raise ValueError("input_tensor_shapes must match columns")

    return multi_inputs

def _validate_and_transform_single_input(
    batch: pd.DataFrame,
    input_shapes: List[Optional[List[int]]],
    has_tensors: bool,
    has_tuple: bool,
) -> np.ndarray:
    """验证并转换单个输入"""
    # 处理逻辑省略（与原文一致）
    return single_input

def _validate_and_transform_prediction_result(
    preds: Union[np.ndarray, Mapping[str, np.ndarray], List[Mapping[str, Any]]],
    num_input_rows: int,
    return_type: DataType,
) -> Union[pd.DataFrame, pd.Series]:
    """验证并转换预测结果"""
    # 处理逻辑省略（与原文一致）
    return pd.DataFrame(preds)

def predict_batch_udf(
    make_predict_fn: Callable[
        [],
        PredictBatchFunction,
    ],
    *,
    return_type: DataType,
    batch_size: int,
    input_tensor_shapes: Optional[Union[List[Optional[List[int]]], Mapping[int, List[int]]]] = None,
):
    """定义批量预测的 Pandas UDF"""
    model_uuid = uuid.uuid4()
    def predict(data: Iterator[Union[pd.Series, pd.DataFrame]]) -> Iterator[pd.DataFrame]:
        from model_cache import ModelCache

        predict_fn = ModelCache.get(model_uuid)
        if not predict_fn:
            predict_fn = make_predict_fn()
            ModelCache.add(model_uuid, predict_fn)

        signature = inspect.signature(predict_fn)
        num_expected_cols = len(signature.parameters)

        input_shapes: List[Optional[List[int]]]
        if isinstance(input_tensor_shapes, Mapping):
            input_shapes = [None] * num_expected_cols
            for index, shape in input_tensor_shapes.items():
                input_shapes[index] = shape
        else:
            input_shapes = input_tensor_shapes  # type: ignore

        for pandas_batch in data:
            has_tuple = isinstance(pandas_batch, Tuple)  # type: ignore
            has_tensors = _has_tensor_cols(pandas_batch)

            if has_tensors and not input_shapes:
                raise ValueError("Tensor columns require input_tensor_shapes")

            for batch in _batched(pandas_batch, batch_size):
                num_input_rows = len(batch)
                num_input_cols = len(batch.columns)
                if num_input_cols == num_expected_cols and num_expected_cols > 1:
                    multi_inputs = _validate_and_transform_multiple_inputs(
                        batch, input_shapes, num_input_cols
                    )
                    preds = predict_fn(*multi_inputs)
                elif num_expected_cols == 1:
                    single_input = _validate_and_transform_single_input(
                        batch, input_shapes, has_tensors, has_tuple
                    )
                    preds = predict_fn(single_input)
                else:
                    msg = "Model expected {} inputs, but received {} columns"
                    raise ValueError(msg.format(num_expected_cols, num_input_cols))

                yield _validate_and_transform_prediction_result(
                    preds, num_input_rows, return_type
                )  # type: ignore

    return pandas_udf(predict, return_type)  # type: ignore[call-overload]

def extract_text_embedding(model, tokenizer, sentence):
    """提取文本嵌入向量"""
    inputs = tokenizer(sentence, return_tensors='pt', max_length=32, padding=True, truncation=True)
    embeddings = model(**inputs)
    embeddings = embeddings.pooler_output
    embeddings = embeddings.tolist()
    for i in range(len(embeddings)):
        embeddings[i] = [round(c,4) for c in embeddings[i]]
    return np.array(embeddings, dtype=np.float32)

if __name__ == "__main__":
    args = parse_args()
    spark = init_spark()    
	
	  ### 读取数据
	  df = spark.sql(f"""
         select article_id, title
         from xxx
     """)
		
    def predict_embedding():
        system_command(f"""rm -rf ./bert-base-chinese""")
        system_command(f"""{hadoop} fs -get /path/to/bert-base-chinese""")
        from transformers import BertTokenizer, BertModel
        tokenizer = BertTokenizer.from_pretrained('./bert-base-chinese')
        text_model = BertModel.from_pretrained('./bert-base-chinese')

        def predict(inputs):
            sentence = inputs.tolist()
            embeddings = extract_text_embedding(text_model, tokenizer, sentence)
            return embeddings
        return predict


    predict_embedding_udf = predict_batch_udf(predict_embedding,
                                  return_type=ArrayType(StringType()),
                                  batch_size=100)


    df.withColumn("title_embedding", predict_embedding_udf("title")).show(5)

    spark.stop()
    del spark

关键点说明：

1. 模型加载与缓存：通过 predict_batch_udf 函数封装预测逻辑，利用模型缓存避免重复加载，提高效率。
2. 批量处理：使用 _batched 函数将数据分批处理，避免内存溢出，适合大数据场景。
3. 类型转换与验证：通过 _validate_and_transform 系列函数确保输入输出类型匹配，提高代码健壮性。
4. Hugging Face 模型集成：在 predict_embedding 函数中加载 Hugging Face 的 BERT 模型，并定义预测逻辑。
5. typing语法修改：python3.10以的typing是不支持｜语法的，需要改成Union进行类型的或推断

方法比较

通过以上两种方法，可以在不同 Spark 版本环境下实现 Hugging Face 和 Spark 的结合，充分发挥两者在 NLP 和大数据处理中的优势，推荐第二种，更加可控一些。

本文由博客一文多发平台 OpenWrite 发布！