Update README.md
Browse files
README.md
CHANGED
@@ -1,3 +1,150 @@
|
|
1 |
-
|
2 |
-
|
3 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# 1 开源清单
|
2 |
+
|
3 |
+
本次开源2个通用向量编码模型和一个针对dialogue进行编码的向量模型,同时开源全量160万对话重写数据集和20万的难负例的检索数据集。
|
4 |
+
|
5 |
+
**开源模型:**
|
6 |
+
|
7 |
+
| ModelName | ModelSize | MaxTokens | EmbeddingDimensions | Language | Scenario | C-MTEB Score |
|
8 |
+
|---------------------------------------------------------------------------------------------------------------|-----------|-----------|---------------------|----------|----------|--------------|
|
9 |
+
| [infgrad/stella-base-zh-v3-1792d](https://huggingface.co/infgrad/stella-base-zh-v3-1792d) | 0.4GB | 512 | 1792 | zh-CN | 通用文本 | 67.96 |
|
10 |
+
| [infgrad/stella-large-zh-v3-1792d](https://huggingface.co/infgrad/stella-large-zh-v3-1792d) | 1.3GB | 512 | 1792 | zh-CN | 通用文本 | 68.48 |
|
11 |
+
| [infgrad/stella-dialogue-large-zh-v3-1792d](https://huggingface.co/infgrad/stella-dialogue-large-zh-v3-1792d) | 1.3GB | 512 | 1792 | zh-CN | **对话文本** | 不适用 |
|
12 |
+
|
13 |
+
**开源数据:**
|
14 |
+
|
15 |
+
1. [全量对话重写数据集](https://huggingface.co/datasets/infgrad/dialogue_rewrite_llm) 约160万
|
16 |
+
2. [部分带有难负例的检索数据集](https://huggingface.co/datasets/infgrad/retrieval_data_llm) 约20万
|
17 |
+
|
18 |
+
上述数据集均使用LLM构造,欢迎各位贡献数据集。
|
19 |
+
|
20 |
+
# 2 使用方法
|
21 |
+
|
22 |
+
## 2.1 通用编码模型使用方法
|
23 |
+
|
24 |
+
直接SentenceTransformer加载即可:
|
25 |
+
|
26 |
+
```python
|
27 |
+
from sentence_transformers import SentenceTransformer
|
28 |
+
|
29 |
+
model = SentenceTransformer("infgrad/stella-base-zh-v3-1792d")
|
30 |
+
# model = SentenceTransformer("infgrad/stella-large-zh-v3-1792d")
|
31 |
+
vectors = model.encode(["text1", "text2"])
|
32 |
+
```
|
33 |
+
|
34 |
+
## 2.2 dialogue编码模型使用方法
|
35 |
+
|
36 |
+
**使用场景:**
|
37 |
+
**在一段对话中,需要根据用户语句去检索相关文本,但是对话中的用户语句存在大量的指代和省略,导致直接使用通用编码模型效果不好,
|
38 |
+
可以使用本项目的专门的dialogue编码模型进行编码**
|
39 |
+
|
40 |
+
**使用要点:**
|
41 |
+
|
42 |
+
1. 对dialogue进行编码时,dialogue中的每个utterance需要是如下格式:`"{ROLE}: {TEXT}"`,然后使用`[SEP]` join一下
|
43 |
+
2. 整个对话都要送入模型进行编码,如果长度不够就删掉早期的对话,**编码后的向量本质是对话中最后一句话的重写版本的向量!!**
|
44 |
+
3. 对话用stella-dialogue-large-zh-v3-1792d编码,被检索文本使用stella-large-zh-v3-1792d进行编码,所以本场景是需要2个编码模型的
|
45 |
+
|
46 |
+
如果对使用方法还有疑惑,请到下面章节阅读该模型是如何训练的。
|
47 |
+
|
48 |
+
使用示例:
|
49 |
+
|
50 |
+
```python
|
51 |
+
from sentence_transformers import SentenceTransformer
|
52 |
+
|
53 |
+
dial_model = SentenceTransformer("infgrad/stella-dialogue-large-zh-v3-1792d")
|
54 |
+
general_model = SentenceTransformer("infgrad/stella-large-zh-v3-1792d")
|
55 |
+
# dialogue = ["张三: 吃饭吗", "李四: 等会去"]
|
56 |
+
dialogue = ["A: 最近去打篮球了吗", "B: 没有"]
|
57 |
+
corpus = ["B没打篮球是因为受伤了。", "B没有打乒乓球"]
|
58 |
+
last_utterance_vector = dial_model.encode(["[SEP]".join(dialogue)], normalize_embeddings=True)
|
59 |
+
corpus_vectors = general_model.encode(corpus, normalize_embeddings=True)
|
60 |
+
# 计算相似度
|
61 |
+
sims = (last_utterance_vector * corpus_vectors).sum(axis=1)
|
62 |
+
print(sims)
|
63 |
+
```
|
64 |
+
|
65 |
+
# 3 通用编码模型训练技巧分享
|
66 |
+
|
67 |
+
## hard negative
|
68 |
+
|
69 |
+
难负例挖掘也是个经典的trick了,几乎总能提升效果
|
70 |
+
|
71 |
+
## dropout-1d
|
72 |
+
|
73 |
+
dropout已经是深度学习的标配,我们可以稍微改造下使其更适合句向量的训练。
|
74 |
+
我们在训练时会尝试让每一个token-embedding都可以表征整个句子,而在推理时使用mean_pooling从而达到类似模型融合的效果。
|
75 |
+
具体操作是在mean_pooling时加入dropout_1d,torch代码如下:
|
76 |
+
|
77 |
+
```python
|
78 |
+
vector_dropout = nn.Dropout1d(0.3) # 算力有限,试了0.3和0.5 两个参数,其中0.3更优
|
79 |
+
last_hidden_state = bert_model(...)[0]
|
80 |
+
last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
|
81 |
+
last_hidden = vector_dropout(last_hidden)
|
82 |
+
vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
|
83 |
+
```
|
84 |
+
|
85 |
+
# 4 dialogue编码模型细节
|
86 |
+
|
87 |
+
## 4.1 为什么需要一个dialogue编码模型?
|
88 |
+
|
89 |
+
参见本人历史文章:https://www.zhihu.com/pin/1674913544847077376
|
90 |
+
|
91 |
+
## 4.2 训练数据
|
92 |
+
|
93 |
+
单条数据示例:
|
94 |
+
|
95 |
+
```json
|
96 |
+
{
|
97 |
+
"dialogue": [
|
98 |
+
"A: 最近去打篮球了吗",
|
99 |
+
"B: 没有"
|
100 |
+
],
|
101 |
+
"last_utterance_rewrite": "B: 我最近没有去打篮球"
|
102 |
+
}
|
103 |
+
```
|
104 |
+
|
105 |
+
## 4.3 训练Loss
|
106 |
+
|
107 |
+
```
|
108 |
+
loss = cosine_loss( dial_model.encode(dialogue), existing_model.encode(last_utterance_rewrite) )
|
109 |
+
```
|
110 |
+
|
111 |
+
dial_model就是要被训练的模型,本人是以stella-large-zh-v3-1792d作为base-model进行继续训练的
|
112 |
+
|
113 |
+
existing_model就是现有训练好的**通用编码模型**,本人使用的是stella-large-zh-v3-1792d
|
114 |
+
|
115 |
+
已开源dialogue-embedding的全量训练数据,理论上可以复现本模型效果。
|
116 |
+
|
117 |
+
Loss下降情况:
|
118 |
+
|
119 |
+
<div align="center">
|
120 |
+
<img src="dial_loss.png" alt="icon" width="2000px"/>
|
121 |
+
</div>
|
122 |
+
|
123 |
+
## 4.4 效果
|
124 |
+
|
125 |
+
目前还没有专门测试集,本人简单测试了下是有效果的,部分测试结果见文件`dial_retrieval_test.xlsx`。
|
126 |
+
|
127 |
+
# 5 后续TODO
|
128 |
+
|
129 |
+
1. 更多的dial-rewrite数据
|
130 |
+
2. 不同EmbeddingDimensions的编码模型
|
131 |
+
|
132 |
+
# 6 FAQ
|
133 |
+
|
134 |
+
Q: 为什么向量维度是1792?\
|
135 |
+
A: 最初考虑发布768、1024,768+768,1024+1024,1024+768维度,但是时间有限,先做了1792就只发布1792维度的模型。理论上维度越高效果越好。
|
136 |
+
|
137 |
+
Q: 如何复现CMTEB效果?\
|
138 |
+
A: SentenceTransformer加载后直接用官方评测脚本就行,注意对于Classification任务向量需要先normalize一下
|
139 |
+
|
140 |
+
Q: 复现的CMTEB效果和本文不一致?\
|
141 |
+
A: 聚类不一致正常,官方评测代码没有设定seed,其他不一致建议检查代码或联系本人。
|
142 |
+
|
143 |
+
Q: 如何选择向量模型?\
|
144 |
+
A: 没有免费的午餐,在自己测试集上试试,本人推荐bge、e5和stella.
|
145 |
+
|
146 |
+
Q: 长度为什么只有512,能否更长?\
|
147 |
+
A: 可以但没必要,长了效果普遍不好,这是当前训练方法和数据导致的,几乎无解,建议长文本还是走分块。
|
148 |
+
|
149 |
+
Q: 训练资源和算力?\
|
150 |
+
A: 亿级别的数据,单卡A100要一个月起步
|