Add citation, delete amp

README.md

@@ -4,7 +4,9 @@ license: apache-2.0
 
 **English** | [中文](./README_zh.md)
 
-## Code implementation of new GTE embeddings
+[Arxiv PDF](https://arxiv.org/pdf/2407.19669), [HF paper page](https://huggingface.co/papers/2407.19669)
+
+## Code implementation of new GTE encoders
 
 This model is a BERT-like encoder with the following optimizations implemented:
 
@@ -12,7 +14,6 @@ This model is a BERT-like encoder with the following optimizations implemented:
 2. Substituting the conventional activation functions with Gated Linear Units (GLU) [^2].
 3. Setting attention dropout to 0 to use `xformers` and `flash_attn`.
 4. Using unpadding to eliminate the needless computations for padding tokens [^3]. (this is off by default and should be used in conjunction with `xformers` for optimal acceleration).
-5. Setting `vocab_size` as a multiple of 64.
 
 ### Recommendation: Enable Unpadding and Acceleration with `xformers`
 
@@ -31,7 +32,8 @@ elif pytorch is installed using pip:
 ```
 For more information, refer to [Installing xformers](https://github.com/facebookresearch/xformers?tab=readme-ov-file#installing-xformers).
 
-Then, when loading the model, set `unpad_inputs` and `use_memory_efficient_attention` to `true`, and enable `fp16` mixed precision computation to achieve the fastest acceleration.
+Then, when loading the model, set `unpad_inputs` and `use_memory_efficient_attention` to `true`,
+and set `torch_dtype` to `torch.float16` (or `torch.bfloat16`) to achieve the acceleration.
 
 ```python
 import torch
@@ -45,15 +47,18 @@ model = AutoModel.from_pretrained(
     trust_remote_code=True,
     unpad_inputs=True,
     use_memory_efficient_attention=True,
+    torch_dtype=torch.float16
 ).to(device)
 
-with torch.autocast(device_type=device.type, dtype=torch.float16):  # or bfloat16
-    with torch.inference_mode():
-        outputs = model(**inputs.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))
 
 ```
 
-Alternatively, you can directly modify the `unpad_inputs` and `use_memory_efficient_attention` settings to `true` in the model's `config.json`, eliminating the need to set them in the code.
+Alternatively, you can directly modify the `unpad_inputs` and `use_memory_efficient_attention` settings to `true` in the model's `config.json`,
+eliminating the need to set them in the code.
 
 
 ---
@@ -73,6 +78,20 @@ Without the outstanding work of `nomicai`, the release of `gte-v1.5` could have
 
 ---
 
+## 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}, 
+}
+```
+
+
 [^1]: Su, Jianlin, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. "Roformer: Enhanced transformer with rotary position embedding." Neurocomputing 568 (2024): 127063.
 
 [^2]: Shazeer, Noam. "Glu variants improve transformer." arXiv preprint arXiv:2002.05202 (2020).

README_zh.md

@@ -4,6 +4,8 @@ license: apache-2.0
 
 [English](./README.md) | **中文**
 
+[Arxiv PDF](https://arxiv.org/pdf/2407.19669), [HF paper page](https://huggingface.co/papers/2407.19669)
+
 ## GTE 新模型代码实现
 
 此模型为 BERT-like 编码器模型，加入了以下优化：
@@ -12,7 +14,6 @@ license: apache-2.0
 2. 使用 GLU (Gated Linear Unit) [^2] 替换普通的激活函数。
 3. 设置 attention dropout 为 0 以方便应用 `xformers` 和 `flash_attn` 等优化。
 4. 使用 Unpadding 技术去除对 padding token 的无用计算 [^3]（默认关闭，需要结合 `flash_attn` 或 `xformers` 使用来获得最高加速）。
-5. 设置 `vocab_size % 64 = 0`。
 
 
 ### 推荐：启用 Unpadding 和 xformers 加速
@@ -32,7 +33,7 @@ elif pytorch 使用 pip 安装 :
 ```
 更多信息可参考 [installing-xformers](https://github.com/facebookresearch/xformers?tab=readme-ov-file#installing-xformers)。
 
-然后，加载模型时设置 `unpad_inputs` 和 `use_memory_efficient_attention` 为 `true`，并启用 `fp16` 混合精度计算，即可获得最快加速。
+然后，加载模型时设置 `unpad_inputs` 和 `use_memory_efficient_attention` 为 `true`，并设置 `torch_dtype` 为 `torch.float16` (or `torch.bfloat16`)，即可获得加速。
 
 ```python
 import torch
@@ -46,11 +47,13 @@ model = AutoModel.from_pretrained(
     trust_remote_code=True,
     unpad_inputs=True,
     use_memory_efficient_attention=True,
+    torch_dtype=torch.float16
 ).to(device)
 
-with torch.autocast(device_type=device.type, dtype=torch.float16):  # 或bfloat16
-    with torch.inference_mode():
-        outputs = model(**inputs.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`，省去代码中的设置。
@@ -72,6 +75,21 @@ with torch.autocast(device_type=device.type, dtype=torch.float16):  # 或bfloat1
 
 ---
 
+
+## 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}, 
+}
+```
+
+
 [^1]: Su, Jianlin, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. "Roformer: Enhanced transformer with rotary position embedding." Neurocomputing 568 (2024): 127063.
 
 [^2]: Shazeer, Noam. "Glu variants improve transformer." arXiv preprint arXiv:2002.05202 (2020).