Introduce code from Rust monorepo. This is a mess

tts/NeMo/examples/nlp/information_retrieval/bert_dpr.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import BertDPRModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="bert_ir_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    bert_dpr_model = BertDPRModel(cfg.model, trainer=trainer)
+    trainer.fit(bert_dpr_model)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/information_retrieval/bert_joint_ir.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import BertJointIRModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+
+
+@hydra_runner(config_path="conf", config_name="bert_ir_config")
+def main(cfg: DictConfig) -> None:
+    logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+    trainer = pl.Trainer(**cfg.trainer)
+    exp_manager(trainer, cfg.get("exp_manager", None))
+    bert_joint_ir_model = BertJointIRModel(cfg.model, trainer=trainer)
+    trainer.fit(bert_joint_ir_model)
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/information_retrieval/conf/bert_ir_config.yaml

@@ -0,0 +1,99 @@
+# Fine-tuning BERT model for information retrieval
+name: &name BertIR
+trainer:
+  gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices
+  num_nodes: 1
+  max_epochs: 2 # the number of training epochs
+  max_steps: null # precedence over max_epochs
+  accumulate_grad_batches: 1 # accumulates grads every k batches
+  precision: 16 # 16 to use AMP
+  amp_level: O1 # O1 or O2 if using AMP
+  accelerator: ddp
+  log_every_n_steps: 1  # Interval of logging.
+  val_check_interval: 0.05  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
+  checkpoint_callback: false  # provided by exp_manager
+  logger: false  # provided by exp_manager
+
+model:
+  nemo_path: null # exported .nemo path
+
+  language_model:
+    pretrained_model_name: bert-base-uncased
+    sim_score_dropout: 0.1
+    lm_checkpoint: null
+    config:
+      attention_probs_dropout_prob: 0.1
+      hidden_act: gelu
+      hidden_dropout_prob: 0.1
+      hidden_size: 768
+      initializer_range: 0.02
+      intermediate_size: 3072
+      max_position_embeddings: 512
+      num_attention_heads: 12
+      num_hidden_layers: 12
+      type_vocab_size: 2
+      vocab_size: 30522
+    config_file: null # json file, precedence over config
+
+  tokenizer:
+    tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec
+    vocab_file: null # path to vocab file
+    tokenizer_model: null # tokenizer model for sentencepiece
+    special_tokens: null
+
+  train_ds:
+    passages: null # path to file with passages and their indices
+    queries: null # path to file with training queries and their indices
+    query_to_passages: null
+    # path to file with training examples which have the form of
+    # (query_id, relevant_passage_id, irrelevant_passage_1_id, ..., irrelevant_passage_n_id)
+    num_negatives: 10
+    batch_size: 6
+    psg_cache_format: npz
+    shuffle: true
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 1
+    drop_last: false
+    pin_memory: false
+
+  validation_ds:
+    passages: null # path to file with passages and their indices
+    queries: null # path to file with validation queries and their indices
+    query_to_passages: null # path to file with passages to re-rank for each validation query
+    num_negatives: 10
+    batch_size: 6
+    psg_cache_format: pkl
+    shuffle: false
+    num_samples: -1 # number of samples to be considered, -1 means all the dataset
+    num_workers: 1
+    drop_last: false
+    pin_memory: false
+
+  optim:
+    name: adam
+    lr: 1e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0
+
+    sched:
+      name: WarmupAnnealing
+      warmup_steps: null
+      warmup_ratio: 0.05
+      last_epoch: -1
+
+      # pytorch lightning args
+      monitor: val_loss
+      reduce_on_plateau: false
+
+exp_manager:
+  exp_dir: null  # where to store logs and checkpoints
+  name: *name  # name of experiment
+  create_tensorboard_logger: True
+  create_checkpoint_callback: True
+
+hydra:
+  run:
+    dir: .
+  job_logging:
+    root:
+      handlers: null

tts/NeMo/examples/nlp/information_retrieval/construct_random_negatives.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+
+import numpy as np
+import pandas as pd
+
+
+def construct_negatives(input_file, output_file, num_passages, num_negatives):
+    qrels = pd.read_csv(input_file, delimiter="\t", header=None)
+    with open(output_file, "w") as f:
+        for i in range(len(qrels)):
+            query_id, rel_passage_id = qrels[0][i], qrels[2][i]
+            negatives = np.random.randint(num_passages, size=num_negatives)
+            output_ids = [query_id, rel_passage_id] + negatives.tolist()
+            output_str = [str(id_) for id_ in output_ids]
+            print("\t".join(output_str), file=f)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Negative passages construction")
+    parser.add_argument("--data", type=str, default="msmarco_dataset", help="path to folder with data")
+    parser.add_argument("--num_passages", type=int, default=8841823, help="total number of passages")
+    parser.add_argument("--num_negatives", type=int, default=10, help="number of negatives per positive")
+    args = parser.parse_args()
+
+    for mode in ["train", "dev"]:
+        construct_negatives(
+            input_file=f"{args.data}/qrels.{mode}.tsv",
+            output_file=f"{args.data}/query2passages.{mode}.tsv",
+            num_passages=args.num_passages,
+            num_negatives=args.num_negatives,
+        )
+
+
+if __name__ == '__main__':
+    main()

tts/NeMo/examples/nlp/information_retrieval/get_msmarco.sh

@@ -0,0 +1,22 @@
+echo "=== Acquiring MSMARCO dataset ==="
+echo "---"
+
+mkdir -p msmarco_dataset
+cd msmarco_dataset
+
+echo "- Downloading passages"
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/collection.tar.gz
+tar -xzvf collection.tar.gz
+rm collection.tar.gz
+
+echo "- Downloading queries"
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/queries.tar.gz
+tar -xzvf queries.tar.gz
+rm queries.tar.gz
+rm queries.eval.tsv
+
+echo "- Downloading relevance labels"
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/qrels.train.tsv
+wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/qrels.dev.tsv
+
+echo "---"