代码详解：distilbert-multilingual-nli-stsb-quora-ranking推理脚本的每一行

代码详解：distilbert-multilingual-nli-stsb-quora-ranking推理脚本的每一行

【免费下载链接】distilbert-multilingual-nli-stsb-quora-ranking项目地址: https://ai.gitcode.com/hf_mirrors/zhouhui/distilbert-multilingual-nli-stsb-quora-ranking

distilbert-multilingual-nli-stsb-quora-ranking是一款强大的多语言句子嵌入模型，能够将不同语言的文本转换为具有语义意义的向量表示。本文将逐行解析其推理脚本examples/inference.py，帮助新手理解模型推理的完整流程。

1. 导入核心依赖库

from openmind import AutoTokenizer, AutoModel, is_torch_npu_available from openmind_hub import snapshot_download import torch.nn.functional as F import torch import argparse

这部分代码导入了模型运行所需的核心库：

• AutoTokenizer和AutoModel：用于自动加载预训练模型和分词器
• is_torch_npu_available：检查是否有NPU加速设备
• torch相关库：提供深度学习计算支持
• argparse：用于解析命令行参数

2. 命令行参数解析

def parse_args(): parser = argparse.ArgumentParser() parser.add_argument( "--model_name_or_path", type=str, help="Path to model", default="zhouhui/distilbert-multilingual-nli-stsb-quora-ranking", ) args = parser.parse_args() return args

parse_args函数定义了一个命令行参数--model_name_or_path，用于指定模型路径，默认值为项目模型名称。这使得用户可以灵活指定不同的模型路径进行推理。

3. 均值池化函数实现

def mean_pooling(model_output, attention_mask): token_embeddings = model_output[0] # First element contains all token embeddings input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

mean_pooling函数是将token级别嵌入转换为句子级别嵌入的关键步骤：

1. 从模型输出中获取token嵌入（model_output[0]）
2. 扩展注意力掩码以匹配token嵌入维度
3. 通过掩码加权平均计算句子嵌入，避免填充token影响结果

4. 主函数实现

4.1 参数解析与设备配置

def main(): args = parse_args() model_path = args.model_name_or_path if is_torch_npu_available(): device = "npu:0" else: device = "cpu"

主函数首先解析命令行参数，然后根据系统环境选择合适的计算设备（NPU或CPU）。

4.2 加载模型与分词器

tokenizer = AutoTokenizer.from_pretrained(model_path) model = AutoModel.from_pretrained(model_path)

这两行代码是加载预训练模型的核心步骤：

• AutoTokenizer.from_pretrained：加载与模型匹配的分词器
• AutoModel.from_pretrained：加载预训练模型权重

4.3 输入文本处理

sentences = ['This is an example sentence', 'Each sentence is converted'] # Tokenize sentences encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')

这里定义了示例句子，并使用分词器对其进行处理：

• padding=True：自动填充到相同长度
• truncation=True：超过最大长度时截断
• return_tensors='pt'：返回PyTorch张量格式

4.4 模型推理与嵌入计算

# Compute token embeddings with torch.no_grad(): model_output = model(**encoded_input) # Perform pooling sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])

推理过程分为两步：

1. 使用torch.no_grad()禁用梯度计算，提高推理速度
2. 将编码后的输入传递给模型，获取token级别嵌入
3. 应用均值池化将token嵌入转换为句子嵌入

4.5 输出结果

print("Sentence embeddings:") print(sentence_embeddings)

最后打印计算得到的句子嵌入向量，这些向量可以用于语义相似度计算、文本分类等下游任务。

5. 程序入口

if __name__ == "__main__": main()

这是Python程序的标准入口方式，确保当脚本被直接运行时才执行main函数。

总结

通过对inference.py脚本的逐行解析，我们了解了distilbert-multilingual-nli-stsb-quora-ranking模型从加载到推理的完整流程。这个简洁的脚本展示了如何使用预训练模型将文本转换为向量表示，为各种NLP应用提供基础支持。要使用此模型，只需克隆仓库：git clone https://gitcode.com/hf_mirrors/zhouhui/distilbert-multilingual-nli-stsb-quora-ranking，然后运行推理脚本即可体验多语言句子嵌入的强大功能。

【免费下载链接】distilbert-multilingual-nli-stsb-quora-ranking项目地址: https://ai.gitcode.com/hf_mirrors/zhouhui/distilbert-multilingual-nli-stsb-quora-ranking