upload

README.md

@@ -0,0 +1,82 @@
+---
+language: ja
+datasets:
+- unicamp-dl/mmarco
+widget:
+- text: "Python（パイソン）はインタープリタ型の高水準汎用プログラミング言語である。グイド・ヴァン・ロッサムにより創り出され、1991年に最初にリリースされたPythonの設計哲学は、有意なホワイトスペース(オフサイドルール)の顕著な使用によってコードの可読性を重視している。その言語構成とオブジェクト指向のアプローチは、プログラマが小規模なプロジェクトから大規模なプロジェクトまで、明確で論理的なコードを書くのを支援することを目的としている。"
+
+license: apache-2.0
+---
+
+# doc2query/msmarco-japanese-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-japanese-mt5-base-v1'
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+text = "Python（パイソン）はインタープリタ型の高水準汎用プログラミング言語である。グイド・ヴァン・ロッサムにより創り出され、1991年に最初にリリースされたPythonの設計哲学は、有意なホワイトスペース(オフサイドルール)の顕著な使用によってコードの可読性を重視している。その言語構成とオブジェクト指向のアプローチは、プログラマが小規模なプロジェクトから大規模なプロジェクトまで、明確で論理的なコードを書くのを支援することを目的としている。"
+
+
+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:48460067426c0ff927645e732547a8a58773fdeba6cf1364ece70d76c5f6e33a
+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 japanese

training_args.bin

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