如何在Snowflake UDF中使用已下载的AutoTokenizer？

Snowflake Python UDF无法通过shutil.copy直接访问Stage路径——Stage文件并非挂载在本地文件系统中，必须通过Snowflake官方的资源导入机制或专用文件API读取。

以下两种方案均可实现UDF调用本地上传的BERT中文分词器：

方案1：通过UDF的IMPORTS参数加载文件 #

这是最高效的方式，直接将Stage中的分词器文件作为资源导入UDF运行环境。

步骤1：确认Stage文件结构 #

先执行以下命令验证文件是否正确上传到指定路径：

LIST @my_db.my_schema.my_stage/tokenizer/;

需确保输出包含tokenizer_config.json、special_tokens_map.json、vocab.txt、tokenizer.json四个文件。

步骤2：修改UDF定义 #

CREATE OR REPLACE FUNCTION my_tokenizer_udf(str STRING)
RETURNS ARRAY
LANGUAGE PYTHON
RUNTIME_VERSION = 3.11
HANDLER = 'tokenize_text'
PACKAGES = ('transformers==4.35.2', 'torch')  # 指定兼容版本，避免依赖问题
IMPORTS = (
    '@my_db.my_schema.my_stage/tokenizer/tokenizer_config.json',
    '@my_db.my_schema.my_stage/tokenizer/special_tokens_map.json',
    '@my_db.my_schema.my_stage/tokenizer/vocab.txt',
    '@my_db.my_schema.my_stage/tokenizer/tokenizer.json'
)
AS
$$
import os
from transformers import AutoTokenizer

def tokenize_text(input_string):
    # 创建临时目录存放分词器文件
    tokenizer_dir = '/tmp/tokenizer'
    os.makedirs(tokenizer_dir, exist_ok=True)
    
    # 将导入的文件移动到统一目录（IMPORTS默认把文件放在/tmp/imports/下）
    files = [
        'tokenizer_config.json',
        'special_tokens_map.json',
        'vocab.txt',
        'tokenizer.json'
    ]
    for file in files:
        src = f'/tmp/imports/{file}'
        dst = f'{tokenizer_dir}/{file}'
        if os.path.exists(src):
            os.rename(src, dst)
    
    # 加载分词器并执行分词
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
    result = tokenizer(input_string)
    return result["input_ids"]
$$;

方案2：使用Snowpark文件API读取Stage文件 #

如果不想通过IMPORTS导入，可直接用Snowpark的文件操作API读取Stage文件并写入临时目录：

CREATE OR REPLACE FUNCTION my_tokenizer_udf(str STRING)
RETURNS ARRAY
LANGUAGE PYTHON
RUNTIME_VERSION = 3.11
HANDLER = 'tokenize_text'
PACKAGES = ('transformers==4.35.2', 'snowflake-snowpark-python')
AS
$$
import os
from snowflake.snowpark.files import SnowflakeFile
from transformers import AutoTokenizer

def tokenize_text(input_string):
    tokenizer_dir = '/tmp/tokenizer'
    os.makedirs(tokenizer_dir, exist_ok=True)
    
    # 遍历读取Stage中的每个分词器文件
    files = [
        'tokenizer_config.json',
        'special_tokens_map.json',
        'vocab.txt',
        'tokenizer.json'
    ]
    for file in files:
        stage_path = f'@my_db.my_schema.my_stage/tokenizer/{file}'
        with SnowflakeFile.open(stage_path, 'rb') as f_in:
            with open(f'{tokenizer_dir}/{file}', 'wb') as f_out:
                f_out.write(f_in.read())
    
    # 分词并返回结果
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
    result = tokenizer(input_string)
    return result["input_ids"]
$$;

• 原代码中tokenizer.tokenize()返回的是token字符串列表，而非包含input_ids的字典，需改用tokenizer(input_string)获取完整分词结果
• 若上传的是整个tokenizer目录，可在IMPORTS中用通配符简化配置：@my_db.my_schema.my_stage/tokenizer/*
• 执行UDF前，确保当前角色拥有Stage的读取权限

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