如何将Spark大数据批量输入HuggingFace Pipeline推理并解决序列过长问题

一、解决序列过长错误 #

错误核心是bert-base-uncased模型的最大输入序列长度为512，你的文本tokenize后长度超出了这个限制，可通过以下三种方式解决：

1. 强制截断文本 #

直接截断超过512长度的token序列，保留前512个token进行推理，适合对长文本尾部信息不敏感的场景：

# 方式1：初始化tokenizer时配置
load_tokenizer = AutoTokenizer.from_pretrained(model_name, truncation=True, max_length=512)

# 方式2：调用pipeline时传入参数
my_pipeline(a, truncation=True, max_length=512)

2. 长文本分块聚合结果 #

将长文本切成多个不超过512token的块，分别推理后合并结果，保留全部文本信息：

def process_long_text(text, pipeline, max_len=512):
    tokens = load_tokenizer.encode(text, add_special_tokens=False)
    # 预留[CLS]和[SEP]的位置，每个块实际最大token数为max_len-2
    chunks = [tokens[i:i+max_len-2] for i in range(0, len(tokens), max_len-2)]
    results = []
    for chunk in chunks:
        chunk_text = load_tokenizer.decode(chunk)
        res = pipeline(chunk_text)[0]
        results.append((res['label'], res['score']))
    
    # 合并策略：取概率最高的标签
    label_scores = {}
    for label, score in results:
        label_scores[label] = label_scores.get(label, []) + [score]
    final_label = max(label_scores, key=lambda x: max(label_scores[x]))
    final_score = max(label_scores[final_label])
    return {'label': final_label, 'score': final_score}

# 批量处理示例
processed_results = [process_long_text(text, my_pipeline) for text in a]

3. 更换长文本模型 #

直接使用支持更长序列的模型，比如allenai/longformer-base-4096（支持4096长度）：

MODEL = "allenai/longformer-base-4096"
model_name = MODEL
load_model = AutoModelForSequenceClassification.from_pretrained(model_name)
load_tokenizer = AutoTokenizer.from_pretrained(model_name)

二、批量处理Spark DataFrame的正确方式 #

直接转Pandas会导致内存溢出（大数据量场景），推荐用Spark的Pandas UDF实现分布式并行推理：

1. 实现分布式推理 #

from pyspark.sql.functions import pandas_udf, col
import pandas as pd

# 每个executor仅初始化一次模型
def init_pipeline():
    MODEL = "bert-base-uncased"
    model_name = MODEL + '-text-classification'
    load_model = AutoModelForSequenceClassification.from_pretrained(model_name)
    load_tokenizer = AutoTokenizer.from_pretrained(model_name, truncation=True, max_length=512)
    return pipeline("text-classification", model=load_model, tokenizer=load_tokenizer)

# 定义Pandas UDF
@pandas_udf("struct<label:string, score:double>")
def batch_classify(texts: pd.Series) -> pd.Series:
    pipeline = init_pipeline()
    results = pipeline(texts.tolist(), batch_size=32, truncation=True, max_length=512)
    return pd.Series(results)

# 应用到Spark DataFrame
df_result = df_0.withColumn("classification", batch_classify(col("lines")))

# 展开结果列（可选）
df_final = df_result.select("lines", "classification.label", "classification.score")
df_final.show()

2. 性能优化建议 #

• 开启Spark动态资源分配，根据数据量调整executor数量与内存
• 调整batch_size参数，平衡推理速度与内存占用
• 若使用GPU，需配置Spark GPU资源，利用transformers的GPU加速能力

内容的提问来源于stack exchange，提问作者Ivan Lee