如何优化PySpark+Whisper语音转文本代码的速度与报错问题？

• 基于PySpark、Whisper开发的语音转文本代码可实现核心功能，但运行速度慢且存在多处错误
• 询问是否可通过asyncio提升PySpark性能，寻求代码优化及提速方案

数据结构示例 #

report.csv对应的Spark DataFrame结构示例：

sparkDataFrame = (["John", '25/5/2023', '<a href="https://jolimont.native-infinity.com.br/api/queueRecords/2023/01/02/1/call-crmsimples-2196928.mp3" target="_blank">Baixar</a>', "00:00:35"])

当前代码（已修正基础语法错误） #

from pyspark.sql import SparkSession 
from pyspark.sql.functions import * 
from bs4 import BeautifulSoup 
import pandas as pd 
import wget 
import os 
import findspark 
from sparknlp.base import * 
from pyspark.sql.types import BinaryType, StringType 
from sparknlp.annotator import * 
import whisper 
import aiohttp 
import asyncio 
import nest_asyncio 
import time 

findspark.init()
# 仅初始化一次SparkSession
spark = SparkSession.builder.appName("database").master("local[*]").enableHiveSupport().getOrCreate()

# 定义数据库列名
cols = ['Agente', 'Data/Hora', 'Gravação', 'Atendimento']

# 加载CSV到Pandas DataFrame
csv = pd.read_csv("report.csv", sep=';', usecols=cols)

# 转换并过滤通话时长
csv['Atendimento'] = pd.to_datetime(csv['Atendimento'])
csv = csv.loc[csv['Atendimento'] > '00:00:15']

# 转为Spark DataFrame
sparkDataFrame = spark.createDataFrame(csv)

# 下载音频函数
def download_audio(url):
   soup = BeautifulSoup(url, 'html.parser')
   link = soup.a['href']

   # 下载音频
   file = wget.download(link, out="audio.tmp")
   os.rename(file, "audio.mp3")

   # 读取二进制数据
   with open("audio.mp3", "rb") as audio:
      data = audio.read()

   os.remove("audio.mp3")
   return data

# 注册UDF
download_audio_udf = udf(download_audio, BinaryType())
sparkDataFrame = sparkDataFrame.withColumn("Audios", download_audio_udf(sparkDataFrame['Gravação']))

sparkDataFrame.show()
time.sleep(15)

# 异步处理音频函数
async def process_audio(row): 
    async with aiohttp.ClientSession() as session: 
        async with session.get(row) as response:
            data = await response.read()
    
    # 生成唯一临时文件名，避免竞态冲突
    temp_file = f"audio_{hash(row)}.mp3"
    with open(temp_file, "wb") as audio:
        audio.write(data)
    
    # 加载模型并转录
    model = whisper.load_model("small")
    result = model.transcribe(temp_file, language="Portuguese")
    transcription_text = result["text"]

    os.remove(temp_file)
    return transcription_text

# 批量异步处理
async def process_audios(audios):
    tasks = []
    for row in audios:
        task = asyncio.create_task(process_audio(row))
        tasks.append(task)
    results = await asyncio.gather(*tasks)
    return results

# 获取音频URL并处理
urls = sparkDataFrame.select("Gravação").take(3)
audios = [row['Gravação'] for row in urls]

nest_asyncio.apply()
loop = asyncio.get_event_loop()
transcriptions = loop.run_until_complete(process_audios(audios))

for transcription in transcriptions:
    print("Transcription:", transcription)

loop.close()

1. 修复核心错误 #

• 缩进问题：原代码中process_audio、process_audios函数存在缩进错误，已修正
• 文件竞态：多个任务共用audio.mp3会导致文件读写冲突，改为生成唯一临时文件名
• 重复初始化SparkSession：移除重复创建逻辑，避免资源浪费
• HTML解析器大小写：将'HTML.parser'改为'html.parser'，避免解析器加载失败

2. 提升PySpark性能的关键优化 #

（1）跳过Pandas中间层，直接用Spark读取数据 #

避免Pandas与Spark之间的序列化/反序列化开销：

# 替换原Pandas读取CSV的代码
sparkDataFrame = spark.read.csv("report.csv", sep=';', header=True, usecols=cols)
# 转换时长列并过滤
sparkDataFrame = sparkDataFrame.withColumn("Atendimento", to_timestamp("Atendimento")) \
                               .filter(col("Atendimento") > to_timestamp(lit("00:00:15")))

（2）用Pandas UDF替代普通Python UDF #

普通Python UDF在Spark中性能极低，Pandas UDF采用向量化处理，大幅提升效率：

from pyspark.sql.functions import pandas_udf
import requests

# 加载Whisper模型一次，避免重复加载
model = whisper.load_model("small", device="cuda" if torch.cuda.is_available() else "cpu")

# 批量下载音频的Pandas UDF
@pandas_udf(BinaryType())
def download_audio_pandas(urls: pd.Series) -> pd.Series:
    def download_single(url):
        soup = BeautifulSoup(url, 'html.parser')
        link = soup.a['href']
        return requests.get(link).content
    
    return urls.apply(download_single)

# 批量转录音频的Pandas UDF
@pandas_udf(StringType())
def transcribe_audio_pandas(audios: pd.Series) -> pd.Series:
    def transcribe_single(audio_bytes):
        # 直接传入二进制数据，跳过文件IO
        result = model.transcribe(audio_bytes, language="Portuguese")
        return result["text"]
    
    return audios.apply(transcribe_single)

# 应用UDF
sparkDataFrame = sparkDataFrame.withColumn("Audios", download_audio_pandas(col("Gravação")))
sparkDataFrame = sparkDataFrame.withColumn("Transcription", transcribe_audio_pandas(col("Audios")))

3. asyncio与Spark的结合方式 #

asyncio无法直接提升Spark的分布式计算性能，但可优化单节点内的IO密集型操作（如音频下载）：

import asyncio
import aiohttp

async def async_download(session, url):
    soup = BeautifulSoup(url, 'html.parser')
    link = soup.a['href']
    async with session.get(link) as response:
        return await response.read()

@pandas_udf(BinaryType())
def download_audio_async(urls: pd.Series) -> pd.Series:
    async def batch_download():
        async with aiohttp.ClientSession() as session:
            tasks = [async_download(session, url) for url in urls]
            return await asyncio.gather(*tasks)
    
    loop = asyncio.get_event_loop()
    results = loop.run_until_complete(batch_download())
    return pd.Series(results)

注：此优化仅适用于单节点IO场景，分布式环境下需依赖Spark的集群资源调度，asyncio无法替代Spark的并行计算能力。

4. 其他提速建议 #

• Whisper模型优化：
  • 有GPU时指定device="cuda"，利用GPU加速转录
  • 改用更小的模型（如"base"）或量化模型（compute_type="int8"）减少计算量
• Spark资源配置：
  • 调整master("local[*]")核心数，集群环境下配置合适的executor内存、核心数
  • 开启Kryo序列化：spark.conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
• 数据分区：对Spark DataFrame按Agente或Data/Hora分区，避免数据倾斜

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