【限时免费】有手就会！gte-small模型本地部署与首次推理全流程实战-优快云博客

有手就会！gte-small模型本地部署与首次推理全流程实战

【免费下载链接】gte-small 项目地址: https://gitcode.com/mirrors/supabase/gte-small

写在前面：硬件门槛

在开始之前，请确保你的设备满足以下最低硬件要求：

推理（Inference）：至少需要4GB内存，建议使用支持CUDA的GPU（如NVIDIA显卡）以加速计算。
微调（Fine-tuning）：需要更高的计算资源，建议使用16GB以上内存和性能较强的GPU。

如果你的设备满足以上要求，恭喜你，可以继续往下看啦！

环境准备清单

在开始安装和运行gte-small模型之前，你需要准备好以下环境：

Python环境：推荐使用Python 3.8或更高版本。
PyTorch：安装支持CUDA的PyTorch版本（如果使用GPU）。
Transformers库：用于加载和运行模型。
Sentence-Transformers库（可选）：如果你计划使用sentence-transformers库。

安装步骤

安装Python：
- 从Python官网下载并安装最新版本的Python。
- 安装完成后，运行以下命令检查版本：
```
python --version
```
安装PyTorch：
- 访问PyTorch官网，根据你的设备选择适合的安装命令。
- 例如，使用CUDA 11.7的安装命令：
```
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
```
安装Transformers库：
```
pip install transformers
```
安装Sentence-Transformers库（可选）：
```
pip install sentence-transformers
```

模型资源获取

gte-small模型可以通过以下方式获取：

使用transformers库直接加载：

from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("Supabase/gte-small")
model = AutoModel.from_pretrained("Supabase/gte-small")

使用sentence-transformers库加载：

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('Supabase/gte-small')

逐行解析“Hello World”代码

以下是一个完整的Python示例代码，我们将逐行解析：

import torch.nn.functional as F
from torch import Tensor
from transformers import AutoTokenizer, AutoModel

def average_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

input_texts = [
    "what is the capital of China?",
    "how to implement quick sort in python?",
    "Beijing",
    "sorting algorithms"
]

tokenizer = AutoTokenizer.from_pretrained("Supabase/gte-small")
model = AutoModel.from_pretrained("Supabase/gte-small")

# Tokenize the input texts
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

# (Optionally) normalize embeddings
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:1] @ embeddings[1:].T) * 100
print(scores.tolist())

代码解析

导入库：
- torch.nn.functional：提供常用的神经网络函数。
- torch.Tensor：PyTorch的张量类。
- transformers：用于加载预训练模型和分词器。
定义average_pool函数：
- 该函数用于对隐藏状态进行平均池化，生成句子的嵌入向量。
- masked_fill：将无效位置（未被注意力掩码覆盖的位置）填充为0。
- sum：对有效位置求和并除以有效位置数量，得到平均值。
输入文本：
- 定义了一个包含4个句子的列表input_texts。
加载模型和分词器：
- 使用AutoTokenizer.from_pretrained加载分词器。
- 使用AutoModel.from_pretrained加载模型。
分词处理：
- tokenizer对输入文本进行分词，生成batch_dict，包含input_ids和attention_mask。
模型推理：
- 将batch_dict输入模型，得到outputs。
- 使用average_pool函数生成嵌入向量。
归一化（可选）：
- 使用F.normalize对嵌入向量进行归一化。
计算相似度：
- 计算第一个句子与其他句子的相似度得分，并打印结果。

运行与结果展示

运行上述代码后，你将看到如下输出（具体数值可能略有不同）：

[[85.123, 12.456, 5.789]]

第一个数值表示第一个句子与第二个句子的相似度（85.123）。
第二个数值表示第一个句子与第三个句子的相似度（12.456）。
第三个数值表示第一个句子与第四个句子的相似度（5.789）。

常见问题（FAQ）与解决方案

1. 模型加载失败

问题：AutoModel.from_pretrained报错。
解决方案：检查网络连接，确保可以访问模型资源。

2. CUDA内存不足

问题：运行时报错CUDA out of memory。
解决方案：减少batch_size或使用CPU运行。

3. 分词器报错

问题：tokenizer无法处理某些字符。
解决方案：确保输入文本为英文，避免特殊字符。

4. 运行速度慢

问题：模型推理速度慢。
解决方案：使用GPU加速或优化代码。

希望这篇教程能帮助你顺利运行gte-small模型！如果有其他问题，欢迎在评论区交流。