Upload 4 files

README.md

@@ -1,3 +1,83 @@
 ---
 license: apache-2.0
+pipeline_tag: sentence-similarity
+tags:
+- cross-lingual
+- multilingual
+- question-answering
+- retrieval
+- sentence-similarity
+- variational
 ---
+
+# VMSST
+
+Published as a long paper at ACL 2023.
+
+Contrastive learning has been successfully used for retrieval of semantically aligned sentences, but it often requires large batch sizes and carefully engineered heuristics to work well. In this paper, we instead propose a generative model for learning multilingual text embeddings which can be used to retrieve or score sentence pairs. Our model operates on parallel data in N languages and, through an approximation we introduce, efficiently encourages source separation in this multilingual setting, separating semantic information that is shared between translations from stylistic or language-specific variation. We show careful large-scale comparisons between contrastive and generation-based approaches for learning multilingual text embeddings, a comparison that has not been done to the best of our knowledge despite the popularity of these approaches. We evaluate this method on a suite of tasks including semantic similarity, bitext mining, and cross-lingual question retrieval––the last of which we introduce in this paper. Overall, our Variational Multilingual Source-Separation Transformer (VMSST) model outperforms both a strong contrastive and generative baseline on these tasks.
+
+## Checkpoints
+
+T5X (Jax): https://storage.googleapis.com/gresearch/vmsst/vmsst-large-2048-t5x.zip
+
+PyTorch: https://storage.googleapis.com/gresearch/vmsst/vmsst-large-2048-pytorch.zip
+
+## Usage
+
+### Installation
+
+1. Clone the following repository from Google Research.
+
+  ```
+  git clone -b master --single-branch https://github.com/google-research/google-research.git
+  ```
+
+2. Make sure `google-research` is the current directory:
+
+  ```
+  cd google-research/vmsst
+  ```
+
+3. Create and activate a new virtualenv:
+
+  ```
+  python -m venv vmsst
+  source vmsst/bin/activate
+  ```
+  
+4. This repository is tested on Python 3.10+. Install required packages:
+
+  ```
+  pip install -r requirements.txt
+  ```
+
+### Test
+
+To test that the checkpoint and installation are working as intended, run:
+
+    bash run.sh
+
+The expected cosine similarity scores for the three sentences pairs are:
+
+    0.2573888301849365,	0.1563197821378708, and 0.28531330823898315.
+    
+### Inference
+
+To embed a list of sentences:
+
+    python score_sentence_pairs.py --sentence_pair_file test_data/test_sentence_pairs.tsv
+
+To score a list of sentence pairs:
+
+    python embed_sentences.py --sentence_file test_data/test_sentences.txt
+
+## Citation
+
+If you use our code or models your work please cite:
+
+    @article{wieting2022beyond,
+      title={Beyond Contrastive Learning: A Variational Generative Model for Multilingual Retrieval},
+      author={Wieting, John and Clark, Jonathan H and Cohen, William W and Neubig, Graham and Berg-Kirkpatrick, Taylor},
+      journal={arXiv preprint arXiv:2212.10726},
+      year={2022}
+    }

config.json

@@ -0,0 +1,35 @@
+{
+  "_name_or_path": "jwieting/vmsst",
+  "architectures": [
+    "MT5EncoderWithProjection"
+  ],
+  "auto_map": {
+    "AutoModel": "jwieting/vmsst--modeling_vmsst.MT5EncoderWithProjection"
+  },
+  "d_ff": 2816,
+  "d_kv": 64,
+  "d_model": 1024,
+  "decoder_start_token_id": 0,
+  "dense_act_fn": "gelu_new",
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "is_gated_act": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "mt5",
+  "num_decoder_layers": 24,
+  "num_heads": 16,
+  "num_layers": 24,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 32,
+  "tie_word_embeddings": false,
+  "tokenizer_class": "T5Tokenizer",
+  "torch_dtype": "float32",
+  "transformers_version": "4.30.2",
+  "use_cache": true,
+  "vocab_size": 250112
+}

modeling_vmsst.py

@@ -0,0 +1,19 @@
+import torch
+import tqdm
+from torch import nn
+from transformers import MT5EncoderModel, MT5PreTrainedModel
+
+class MT5EncoderWithProjection(MT5PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.mt5_encoder = MT5EncoderModel(config)
+        self.projection = nn.Linear(config.d_model, config.d_model, bias=False)
+        self.post_init()
+
+    def forward(self, **input_args):
+        hidden_states = self.mt5_encoder(**input_args).last_hidden_state
+        mask = input_args['attention_mask']
+        batch_embeddings = torch.sum(hidden_states * mask[:, :, None], dim=1) / torch.sum(mask, dim=1)[:, None]
+        batch_embeddings = self.projection(batch_embeddings)
+        return batch_embeddings

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed70bb240affcc90a945b5905dc643778806ecf9e3c1ff6542de24fa70056228
+size 2262056637