如何在Palantir代码库中加载数据集并将完整目录传入HuggingFace函数？

前提准备 #

将HuggingFace预训练模型（如dslim/bert-base-NER）的所有相关文件（包括config.json、pytorch_model.bin、tokenizer_config.json、vocab.txt等）完整下载后，上传到Palantir的数据集，记录数据集的路径（示例路径：/Users/model/dslim-bert-base-NER）。

代码实现 #

通过Palantir的transform函数，将数据集内的模型文件复制到本地临时目录，再传入HuggingFace的加载函数：

import tempfile
import shutil
from transformers import AutoTokenizer, AutoModelForTokenClassification
from transforms.api import transform, Input

@transform(
    model_dataset=Input("/Users/model/dslim-bert-base-NER")
)
def load_hf_model(model_dataset):
    # 创建临时目录存储模型文件
    with tempfile.TemporaryDirectory() as temp_dir:
        # 遍历数据集内的所有模型文件
        for file_path in model_dataset.filesystem().glob("*"):
            file_name = file_path.split("/")[-1]
            # 将数据集文件复制到临时目录
            with model_dataset.filesystem().open(file_path, "rb") as src:
                with open(f"{temp_dir}/{file_name}", "wb") as dest:
                    shutil.copyfileobj(src, dest)
        
        # 从临时目录加载tokenizer和模型
        tokenizer = AutoTokenizer.from_pretrained(temp_dir)
        model = AutoModelForTokenClassification.from_pretrained(temp_dir)
    
    return tokenizer, model

注意事项 #

• 确保上传的数据集包含模型所需的全部文件，缺失文件会导致加载失败。
• 临时目录会在函数执行完成后自动销毁，无需手动清理。
• 若模型文件在数据集的子目录下，需调整glob的路径参数（如"model_subdir/*"）以匹配实际结构。

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