README_zh.md

metadata

license: apache-2.0

English | 中文

Qwen2 模型代码实现

此模型代码适用于基于Qwen2 的文本表示模型。

默认启用双向注意力机制。

使用方法

1. 下载此仓库中的 configuration.py 和 modeling.py 到你本地保存的 gte-Qwen2 模型目录
2. 将 config.json 的 auto_map 下所有的 modeling_qwen. 替换为 modeling.

推荐：启用 Unpadding 和 xformers 加速

此代码支持使用 xformers 加速 attention 计算，可以根据设备类型自动选择优化实现，比如 flash_attn。 通过 xformers，在不能支持 flash_attn 的旧设备比如V100上也可以获得极大的加速。

首先，安装 xformers（需要预先安装pytorch）：

if pytorch 使用 conda 安装 :
    conda install xformers -c xformers

elif pytorch 使用 pip 安装 :
    # cuda 11.8 version
    pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu118
    # cuda 12.1 version
    pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu121

更多信息可参考 installing-xformers。

然后，加载模型时设置 unpad_inputs 和 use_memory_efficient_attention 为 true，并设置 torch_dtype 为 torch.float16 (or torch.bfloat16)，即可获得加速。

import torch
from transformers import AutoModel, AutoTokenizer

path = 'Alibaba-NLP/gte-Qwen2-1.5B-instruct'
device = torch.device('cuda')
tokenzier = AutoTokenizer.from_pretrained(path)
model = AutoModel.from_pretrained(
    path,
    trust_remote_code=True,
    unpad_inputs=True,
    use_memory_efficient_attention=True,
    torch_dtype=torch.float16
).to(device)

inputs = tokenzier(['test input'], truncation=True, max_length=8192, padding=True, return_tensors='pt')

with torch.inference_mode():
    outputs = model(**inputs.to(device))

也可以直接修改模型的 config.json 中 unpad_inputs 和 use_memory_efficient_attention 为 true，省去代码中的设置。

Citation

@misc{zhang2024mgte,
  title={mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval}, 
  author={Xin Zhang and Yanzhao Zhang and Dingkun Long and Wen Xie and Ziqi Dai and Jialong Tang and Huan Lin and Baosong Yang and Pengjun Xie and Fei Huang and Meishan Zhang and Wenjie Li and Min Zhang},
  year={2024},
  eprint={2407.19669},
  archivePrefix={arXiv},
  primaryClass={cs.CL},
  url={https://arxiv.org/abs/2407.19669}, 
}