readme: ---
  license: cc-by-4.0
  language: 
  - english
  thumbnail: 
  tags:
  - language model
  ---
  !---
  
  # ##############################################################################################
  # 
  # This model has been uploaded to HuggingFace by https://huggingface.co/drAbreu
  # The model is based on the NVIDIA checkpoint located at 
  # https://catalog.ngc.nvidia.com/orgs/nvidia/models/biomegatron345mcased 
  #
  # ##############################################################################################
  -->
  
  [BioMegatron](https://arxiv.org/pdf/2010.06060.pdf) is a transformer developed by the Applied Deep Learning Research team at NVIDIA. This particular Megatron model trained on top of the Megatron-LM model, adding a PubMed corpusto the Megatron-LM corpora(Wikipedia, RealNews, OpenWebText, and CC-Stories). BioMegatron follows a similar (albeit not identical) architecture as BERT and it has 345 million parameters:
  * 24 layers
  * 16 attention heads with a hidden size of 1024. 
  
  More information available at [nVIDIA NGC CATALOG](https://catalog.ngc.nvidia.com/orgs/nvidia/models/biomegatron345mcased)
  
  
  # Running BioMegatron in 🤗 transformers
  
  In this implementation we have followed the commands of the [`nvidia/megatron-bert-uncased-345m`](https://huggingface.co/nvidia/megatron-bert-cased-345m) repository to make BioMegatron available in 🤗. 
  
  However, the file [`convert_megatron_bert_checkpoint.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py) needed a modification. The reason is that the Megatron model shown in [`nvidia/megatron-bert-uncased-345m`](https://huggingface.co/nvidia/megatron-bert-cased-345m) has included head layers, while the weights of the BioMegatron model that we upload to this repository do not contain a head.
  
  The code below is a modification of the original [`convert_megatron_bert_checkpoint.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py).
  
  ```python
  import os
  import torch
  from convert_biomegatron_checkpoint import convert_megatron_checkpoint
  
  print_checkpoint_structure = True
  path_to_checkpoint = "/path/to/BioMegatron345mUncased/"
  
  # Extract the basename.
  basename = os.path.dirname(path_to_checkpoint).split('/')[-1]
  
  # Load the model.
  input_state_dict = torch.load(os.path.join(path_to_checkpoint, 'model_optim_rng.pt'), map_location="cpu")
  
  # Convert.
  print("Converting")
  output_state_dict, output_config = convert_megatron_checkpoint(input_state_dict, head_model=False)
  
  # Print the structure of converted state dict.
  if print_checkpoint_structure:
      recursive_print(None, output_state_dict)
  
  # Store the config to file.
  output_config_file = os.path.join(path_to_checkpoint, "config.json")
  print(f'Saving config to "{output_config_file}"')
  with open(output_config_file, "w") as f:
      json.dump(output_config, f)
  
  # Store the state_dict to file.
  output_checkpoint_file = os.path.join(path_to_checkpoint, "pytorch_model.bin")
  print(f'Saving checkpoint to "{output_checkpoint_file}"')
  torch.save(output_state_dict, output_checkpoint_file)
  
  ```
  
  We provide in the repository an alternative version of the [python script](https://huggingface.co/EMBO/BioMegatron345mCased/blob/main/convert_biomegatron_checkpoint.py) in order to any user to cross-check the validity of the model replicated in this repository.
  
  BioMegatron can be run with the standard 🤗 script for loading models. Here we show an example identical to that of [`nvidia/megatron-bert-uncased-345m`](https://huggingface.co/nvidia/megatron-bert-cased-345m).
  
  ```python
  import os
  import torch
  from transformers import BertTokenizer, MegatronBertForMaskedLM, AutoModelForMaskedLM
  checkpoint = "EMBO/BioMegatron345mCased"
  # The tokenizer. Megatron was trained with standard tokenizer(s).
  tokenizer = BertTokenizer.from_pretrained(checkpoint)
  # Load the model from $MYDIR/nvidia/megatron-bert-uncased-345m.
  model = AutoModelForMaskedLM.from_pretrained(checkpoint)
  device = torch.device("cpu")
  # Create inputs (from the BERT example page).
  input = tokenizer("The capital of France is [MASK]", return_tensors="pt").to(device)
  label = tokenizer("The capital of France is Paris",  return_tensors="pt")["input_ids"].to(device)
  # Run the model.
  with torch.no_grad():
      output = model(**input, labels=label)
      print(output)
  ```
  
  # Limitations
  This implementation has not been fine-tuned in any task. It has only the weights of the official nVIDIA checkpoint. It needs to be trained to perform any downstream task.
  
  # Original code
  The original code for Megatron can be found here: [https://github.com/NVIDIA/Megatron-LM](https://github.com/NVIDIA/Megatron-LM).
  
  