[Flax] Correct all return tensors to numpy (huggingface#13307)

patrickvonplaten · web-flow · commit 2bef3433e5fa · 2021-08-27T17:38:34.000+02:00
* fix_torch_device_generate_test

* remove @

* finish find and replace
diff --git a/examples/research_projects/jax-projects/big_bird/evaluate.py b/examples/research_projects/jax-projects/big_bird/evaluate.py
@@ -110,7 +110,7 @@ def evaluate(example):
         inputs = tokenizer(
             example["question"],
             example["context"],
-            return_tensors="jax",
+            return_tensors="np",
             max_length=4096,
             padding="max_length",
             truncation=True,
diff --git a/src/transformers/file_utils.py b/src/transformers/file_utils.py
@@ -1121,7 +1121,7 @@ def _prepare_output_docstrings(output_type, config_class):
         >>> tokenizer = {tokenizer_class}.from_pretrained('{checkpoint}')
         >>> model = {model_class}.from_pretrained('{checkpoint}')
 
-        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax")
+        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="np")
         >>> outputs = model(**inputs)
 
         >>> # retrieve logts for next token
diff --git a/src/transformers/generation_flax_utils.py b/src/transformers/generation_flax_utils.py
@@ -231,7 +231,7 @@ def generate(
             >>> model = FlaxAutoModelForCausalLM.from_pretrained("distilgpt2")
             >>> input_context = "The dog"
             >>> # encode input context
-            >>> input_ids = tokenizer(input_context, return_tensors="jax").input_ids
+            >>> input_ids = tokenizer(input_context, return_tensors="np").input_ids
             >>> # generate candidates using sampling
             >>> outputs = model.generate(input_ids=input_ids, max_length=20, top_k=30, do_sample=True)
             >>> print("Generated:", tokenizer.batch_decode(outputs, skip_special_tokens=True))
diff --git a/src/transformers/models/bert/modeling_flax_bert.py b/src/transformers/models/bert/modeling_flax_bert.py
@@ -757,7 +757,7 @@ class FlaxBertForPreTraining(FlaxBertPreTrainedModel):
         >>> tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
         >>> model = FlaxBertForPreTraining.from_pretrained('bert-base-uncased')
 
-        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax")
+        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="np")
         >>> outputs = model(**inputs)
 
         >>> prediction_logits = outputs.prediction_logits
diff --git a/src/transformers/models/big_bird/modeling_flax_big_bird.py b/src/transformers/models/big_bird/modeling_flax_big_bird.py
@@ -1567,7 +1567,7 @@ class FlaxBigBirdForPreTraining(FlaxBigBirdPreTrainedModel):
         >>> tokenizer = BigBirdTokenizer.from_pretrained('google/bigbird-roberta-base')
         >>> model = FlaxBigBirdForPreTraining.from_pretrained('google/bigbird-roberta-base')
 
-        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax")
+        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="np")
         >>> outputs = model(**inputs)
 
         >>> prediction_logits = outputs.prediction_logits
diff --git a/src/transformers/models/electra/modeling_flax_electra.py b/src/transformers/models/electra/modeling_flax_electra.py
@@ -761,7 +761,7 @@ class FlaxElectraForPreTraining(FlaxElectraPreTrainedModel):
         >>> tokenizer = ElectraTokenizer.from_pretrained('google/electra-small-discriminator')
         >>> model = FlaxElectraForPreTraining.from_pretrained('google/electra-small-discriminator')
 
-        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax")
+        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="np")
         >>> outputs = model(**inputs)
 
         >>> prediction_logits = outputs.logits
diff --git a/src/transformers/models/vit/modeling_flax_vit.py b/src/transformers/models/vit/modeling_flax_vit.py
@@ -512,7 +512,7 @@ class FlaxViTModel(FlaxViTPreTrainedModel):
         >>> feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
         >>> model = FlaxViTModel.from_pretrained('google/vit-base-patch16-224-in21k')
 
-        >>> inputs = feature_extractor(images=image, return_tensors="jax")
+        >>> inputs = feature_extractor(images=image, return_tensors="np")
         >>> outputs = model(**inputs)
         >>> last_hidden_states = outputs.last_hidden_state
 """
@@ -592,7 +592,7 @@ class FlaxViTForImageClassification(FlaxViTPreTrainedModel):
         >>> feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
         >>> model = FlaxViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
 
-        >>> inputs = feature_extractor(images=image, return_tensors="jax")
+        >>> inputs = feature_extractor(images=image, return_tensors="np")
         >>> outputs = model(**inputs)
         >>> logits = outputs.logits
 
diff --git a/tests/test_modeling_flax_bart.py b/tests/test_modeling_flax_bart.py
@@ -453,7 +453,7 @@ def test_cnn_summarization_same_as_fairseq(self):
             padding="max_length",
             truncation_strategy="only_first",
             truncation=True,
-            return_tensors="jax",
+            return_tensors="np",
         )
 
         self.assertEqual(1024, dct["input_ids"].shape[1])
diff --git a/tests/test_modeling_flax_gpt2.py b/tests/test_modeling_flax_gpt2.py
@@ -213,7 +213,7 @@ def test_use_cache_forward_with_attn_mask(self):
     @slow
     def test_batch_generation(self):
         tokenizer = GPT2Tokenizer.from_pretrained("gpt2", pad_token="</s>", padding_side="left")
-        inputs = tokenizer(["Hello this is a long string", "Hey"], return_tensors="jax", padding=True, truncation=True)
+        inputs = tokenizer(["Hello this is a long string", "Hey"], return_tensors="np", padding=True, truncation=True)
 
         model = FlaxGPT2LMHeadModel.from_pretrained("gpt2")
         model.do_sample = False
diff --git a/tests/test_modeling_flax_gpt_neo.py b/tests/test_modeling_flax_gpt_neo.py
@@ -204,7 +204,7 @@ def test_use_cache_forward_with_attn_mask(self):
     @slow
     def test_batch_generation(self):
         tokenizer = GPT2Tokenizer.from_pretrained("gpt2", pad_token="<|endoftext|>", padding_side="left")
-        inputs = tokenizer(["Hello this is a long string", "Hey"], return_tensors="jax", padding=True, truncation=True)
+        inputs = tokenizer(["Hello this is a long string", "Hey"], return_tensors="np", padding=True, truncation=True)
 
         model = FlaxGPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
         model.do_sample = False

Original file line number	Diff line number	Diff line change
`@@ -453,7 +453,7 @@ def test_cnn_summarization_same_as_fairseq(self):`
`453`	`453`	`padding="max_length",`
`454`	`454`	`truncation_strategy="only_first",`
`455`	`455`	`truncation=True,`
`456`		`- return_tensors="jax",`
	`456`	`+ return_tensors="np",`
`457`	`457`	`)`
`458`	`458`
`459`	`459`	`self.assertEqual(1024, dct["input_ids"].shape[1])`