Update modeling_retrieva_bert.py

modeling_retrieva_bert.py

@@ -65,7 +65,7 @@ from .configuration_retrieva_bert import RetrievaBertConfig
 logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "RetrievaBertConfig"
-_CHECKPOINT_FOR_DOC = "nvidia/megatron-bert-cased-345m"
+_CHECKPOINT_FOR_DOC = "retrieva-jp/bert-1.3b"
 
 
 def load_tf_weights_in_megatron_bert(model, config, tf_checkpoint_path):
@@ -1170,8 +1170,8 @@ class RetrievaBertForPreTraining(RetrievaBertPreTrainedModel):
         >>> from models import RetrievaBertForPreTraining
         >>> import torch
 
-        >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
-        >>> model = RetrievaBertForPreTraining.from_pretrained("nvidia/megatron-bert-cased-345m")
+        >>> tokenizer = AutoTokenizer.from_pretrained("retrieva-jp/bert-1.3b")
+        >>> model = RetrievaBertForPreTraining.from_pretrained("retrieva-jp/bert-1.3b")
 
         >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
         >>> outputs = model(**inputs)
@@ -1294,8 +1294,8 @@ class RetrievaBertForCausalLM(RetrievaBertPreTrainedModel):
         >>> from models import RetrievaBertForCausalLM, RetrievaBertConfig
         >>> import torch
 
-        >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
-        >>> model = RetrievaBertForCausalLM.from_pretrained("nvidia/megatron-bert-cased-345m", is_decoder=True)
+        >>> tokenizer = AutoTokenizer.from_pretrained("retrieva-jp/bert-1.3b")
+        >>> model = RetrievaBertForCausalLM.from_pretrained("retrieva-jp/bert-1.3b", is_decoder=True)
 
         >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
         >>> outputs = model(**inputs)
@@ -1528,8 +1528,8 @@ class RetrievaBertForNextSentencePrediction(RetrievaBertPreTrainedModel):
         >>> from models import RetrievaBertForNextSentencePrediction
         >>> import torch
 
-        >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
-        >>> model = RetrievaBertForNextSentencePrediction.from_pretrained("nvidia/megatron-bert-cased-345m")
+        >>> tokenizer = AutoTokenizer.from_pretrained("retrieva-jp/bert-1.3b")
+        >>> model = RetrievaBertForNextSentencePrediction.from_pretrained("retrieva-jp/bert-1.3b")
 
         >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
         >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."