File size: 2,725 Bytes
ef0f62d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
---
license: apache-2.0
---

[English](./README.md) | **中文**

<!-- [Arxiv PDF](https://arxiv.org/pdf/2407.19669), [HF paper page](https://huggingface.co/papers/2407.19669) -->

## 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](https://github.com/facebookresearch/xformers?tab=readme-ov-file#installing-xformers)。

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

```python
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}, 
}
```