This is a lightweight implementation of SimCSE. You can easily customize the model based on this implementation, and this is more user-friendly to python beginners. (Only supervised SimCSE is supported in this version)
The supervised SimCSE dataset is ./dataset/nli_for_simcse.csv
. Training supervised SimCSE with and without hard negatives are supported. You can specify this setting in Config.py
.
In this version, we only support BERT as the backbone of SimCSE. In order to correctly load the model, bert-base-uncased model needs to be downloaded into ./model
folder.
First, parameters need to be correctly set in Config.py
. This is an approach to replace the command line parameters. Explanation of the parameters can be found within the file.
Next, you can do the training by simply calling python train.py
In this version, we only support STS12 task here. Prediction command: python prediction.py
Besides, In order to show that this implementation is correct, here's the results comparison of the STS12 task (we train the official SimCSE with the same configuration and test it on STS12 task):
Implementation | Spearman's correlation |
---|---|
Score in paper | 0.753 |
SimCSE official(rerun) | 0.761 |
Our implementation | 0.763 |
Hyperparameters:
epoch: 1
batch_size: 32
temperature: 0.05
pooler: CLS
truncate: 32
- Unsupervised SimCSE.
- Training along with the MLM objective.
- Other STS tasks (and perhaps some transfer tasks).
- Chinese version of SimCSE by pytorch.