upload

README.md

@@ -0,0 +1,82 @@
+---
+language: ru
+datasets:
+- unicamp-dl/mmarco
+widget:
+- text: "Python (МФА: [ˈpʌɪθ(ə)n]; в русском языке встречаются названия пито́н или па́йтон) — высокоуровневый язык программирования общего назначения с динамической строгой типизацией и автоматическим управлением памятью, ориентированный на повышение производительности разработчика, читаемости кода и его качества, а также на обеспечение переносимости написанных на нём программ."
+
+license: apache-2.0
+---
+
+# doc2query/msmarco-russian-mt5-base-v1
+
+This is a [doc2query](https://arxiv.org/abs/1904.08375) model based on mT5 (also known as [docT5query](https://cs.uwaterloo.ca/~jimmylin/publications/Nogueira_Lin_2019_docTTTTTquery-v2.pdf)).
+
+It can be used for:
+- **Document expansion**: You generate for your paragraphs 20-40 queries and index the paragraphs and the generates queries in a standard BM25 index like Elasticsearch, OpenSearch, or Lucene. The generated queries help to close the lexical gap of lexical search, as the generate queries contain synonyms. Further, it re-weights words giving important words a higher weight even if they appear seldomn in a paragraph. In our [BEIR](https://arxiv.org/abs/2104.08663) paper we showed that BM25+docT5query is a powerful search engine. In the [BEIR repository](https://github.com/beir-cellar/beir) we have an example how to use docT5query with Pyserini.
+- **Domain Specific Training Data Generation**: It can be used to generate training data to learn an embedding model. In our [GPL-Paper](https://arxiv.org/abs/2112.07577) / [GPL Example on SBERT.net](https://www.sbert.net/examples/domain_adaptation/README.html#gpl-generative-pseudo-labeling) we have an example how to use the model to generate (query, text) pairs for a given collection of unlabeled texts. These pairs can then be used to train powerful dense embedding models.
+
+## Usage
+```python
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+
+model_name = 'doc2query/msmarco-russian-mt5-base-v1'
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+text = "Python (МФА: [ˈpʌɪθ(ə)n]; в русском языке встречаются названия пито́н или па́йтон) — высокоуровневый язык программирования общего назначения с динамической строгой типизацией и автоматическим управлением памятью, ориентированный на повышение производительности разработчика, читаемости кода и его качества, а также на обеспечение переносимости написанных на нём программ."
+
+
+def create_queries(para):
+    input_ids = tokenizer.encode(para, return_tensors='pt')
+    with torch.no_grad():
+        # Here we use top_k / top_k random sampling. It generates more diverse queries, but of lower quality
+        sampling_outputs = model.generate(
+            input_ids=input_ids,
+            max_length=64,
+            do_sample=True,
+            top_p=0.95,
+            top_k=10, 
+            num_return_sequences=5
+            )
+        
+        # Here we use Beam-search. It generates better quality queries, but with less diversity
+        beam_outputs = model.generate(
+            input_ids=input_ids, 
+            max_length=64, 
+            num_beams=5, 
+            no_repeat_ngram_size=2, 
+            num_return_sequences=5, 
+            early_stopping=True
+        )
+
+
+    print("Paragraph:")
+    print(para)
+    
+    print("\nBeam Outputs:")
+    for i in range(len(beam_outputs)):
+        query = tokenizer.decode(beam_outputs[i], skip_special_tokens=True)
+        print(f'{i + 1}: {query}')
+
+    print("\nSampling Outputs:")
+    for i in range(len(sampling_outputs)):
+        query = tokenizer.decode(sampling_outputs[i], skip_special_tokens=True)
+        print(f'{i + 1}: {query}')
+
+create_queries(text)
+
+```
+
+**Note:** `model.generate()` is non-deterministic for top_k/top_n sampling. It produces different queries each time you run it.
+
+## Training
+This model fine-tuned [google/mt5-base](https://huggingface.co/google/mt5-base) for 66k training steps (4 epochs on the 500k training pairs from MS MARCO). For the  training script, see the `train_script.py` in this repository.
+
+The input-text was truncated to 320 word pieces. Output text was generated up to 64 word pieces. 
+
+This model was trained on a (query, passage) from the [mMARCO dataset](https://github.com/unicamp-dl/mMARCO).
+
+
+

config.json

@@ -0,0 +1,30 @@
+{
+  "_name_or_path": "google/mt5-base",
+  "architectures": [
+    "MT5ForConditionalGeneration"
+  ],
+  "d_ff": 2048,
+  "d_kv": 64,
+  "d_model": 768,
+  "decoder_start_token_id": 0,
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "mt5",
+  "num_decoder_layers": 12,
+  "num_heads": 12,
+  "num_layers": 12,
+  "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.18.0",
+  "use_cache": true,
+  "vocab_size": 250112
+}

pytorch_model.bin

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

special_tokens_map.json

@@ -0,0 +1 @@
+{"eos_token": "</s>", "unk_token": "<unk>", "pad_token": "<pad>"}

spiece.model

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

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1 @@
+{"eos_token": "</s>", "unk_token": "<unk>", "pad_token": "<pad>", "extra_ids": 0, "additional_special_tokens": null, "special_tokens_map_file": "/home/patrick/.cache/torch/transformers/685ac0ca8568ec593a48b61b0a3c272beee9bc194a3c7241d15dcadb5f875e53.f76030f3ec1b96a8199b2593390c610e76ca8028ef3d24680000619ffb646276", "name_or_path": "google/mt5-base", "sp_model_kwargs": {}, "tokenizer_class": "T5Tokenizer"}

train_script.py

@@ -0,0 +1,164 @@
+import argparse
+import logging
+from torch.utils.data import Dataset, IterableDataset
+import gzip
+import json
+from transformers import Seq2SeqTrainer, AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainingArguments
+import sys
+from datetime import datetime
+import torch
+import random
+from shutil import copyfile
+import os
+import wandb
+import random
+import re
+from datasets import load_dataset
+import tqdm
+
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    handlers=[logging.StreamHandler(sys.stdout)],
+)
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--lang", required=True)
+parser.add_argument("--model_name", default="google/mt5-base")
+parser.add_argument("--epochs", default=4, type=int)
+parser.add_argument("--batch_size", default=32, type=int)
+parser.add_argument("--max_source_length", default=320, type=int)
+parser.add_argument("--max_target_length", default=64, type=int)
+parser.add_argument("--eval_size", default=1000, type=int)
+#parser.add_argument("--fp16", default=False, action='store_true')
+args = parser.parse_args()
+
+wandb.init(project="doc2query", name=f"{args.lang}-{args.model_name}")
+
+
+
+
+
+def main():
+    ############ Load dataset
+    queries = {}
+    for row in tqdm.tqdm(load_dataset('unicamp-dl/mmarco', f'queries-{args.lang}')['train']):
+        queries[row['id']] = row['text']
+
+    """
+    collection = {}
+    for row in tqdm.tqdm(load_dataset('unicamp-dl/mmarco', f'collection-{args.lang}')['collection']):
+        collection[row['id']] = row['text']
+    """
+    collection = load_dataset('unicamp-dl/mmarco', f'collection-{args.lang}')['collection']
+
+    train_pairs = []
+    eval_pairs = []
+
+
+    with open('qrels.train.tsv') as fIn:
+        for line in fIn:
+            qid, _, did, _ = line.strip().split("\t")
+
+            qid = int(qid)
+            did = int(did)
+
+            assert did == collection[did]['id']
+            text = collection[did]['text']
+
+            pair = (queries[qid], text)
+            if len(eval_pairs) < args.eval_size:
+                eval_pairs.append(pair)
+            else:
+                train_pairs.append(pair)
+
+  
+    print(f"Train pairs: {len(train_pairs)}")
+
+
+    ############ Model
+    model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name)
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+
+    save_steps = 1000
+
+    output_dir = 'output/'+args.lang+'-'+args.model_name.replace("/", "-")+'-'+datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    print("Output dir:", output_dir)
+
+    # Write self to path
+    os.makedirs(output_dir, exist_ok=True)
+
+    train_script_path = os.path.join(output_dir, 'train_script.py')
+    copyfile(__file__, train_script_path)
+    with open(train_script_path, 'a') as fOut:
+        fOut.write("\n\n# Script was called via:\n#python " + " ".join(sys.argv))
+
+    ####
+
+    training_args = Seq2SeqTrainingArguments(
+        output_dir=output_dir,
+        bf16=True,
+        per_device_train_batch_size=args.batch_size,
+        evaluation_strategy="steps",
+        save_steps=save_steps,
+        logging_steps=100,
+        eval_steps=save_steps, #logging_steps,
+        warmup_steps=1000,
+        save_total_limit=1,
+        num_train_epochs=args.epochs,
+        report_to="wandb",
+    )
+
+    ############ Arguments
+
+    ############ Load datasets
+
+
+    print("Input:", train_pairs[0][1])
+    print("Target:", train_pairs[0][0])
+
+    print("Input:", eval_pairs[0][1])
+    print("Target:", eval_pairs[0][0])
+
+
+    def data_collator(examples):
+        targets = [row[0] for row in examples]
+        inputs = [row[1] for row in examples]
+        label_pad_token_id = -100
+
+        model_inputs = tokenizer(inputs, max_length=args.max_source_length, padding=True, truncation=True, return_tensors='pt', pad_to_multiple_of=8 if training_args.fp16 else None)
+
+        # Setup the tokenizer for targets
+        with tokenizer.as_target_tokenizer():
+            labels = tokenizer(targets, max_length=args.max_target_length, padding=True, truncation=True, pad_to_multiple_of=8 if training_args.fp16 else None)
+
+        # replace all tokenizer.pad_token_id in the labels by -100 to ignore padding in the loss.
+        labels["input_ids"] = [
+            [(l if l != tokenizer.pad_token_id else label_pad_token_id) for l in label] for label in labels["input_ids"]
+        ]
+
+
+        model_inputs["labels"] = torch.tensor(labels["input_ids"])
+        return model_inputs
+
+    ## Define the trainer
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_pairs,
+        eval_dataset=eval_pairs,
+        tokenizer=tokenizer,
+        data_collator=data_collator
+    )
+
+    ### Save the model
+    train_result = trainer.train()
+    trainer.save_model()
+    
+    
+if __name__ == "__main__":
+    main()
+
+# Script was called via:
+#python train_hf_trainer_multilingual.py --lang russian

training_args.bin

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