infgrad commited on
Commit
d16fb52
1 Parent(s): c2b56ff

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +150 -3
README.md CHANGED
@@ -1,3 +1,150 @@
1
- ---
2
- license: mit
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要一个月起步