From e5fcc2333959e0fcea7afa49b1c0a18e785f4427 Mon Sep 17 00:00:00 2001
From: fabiocapsouza <fabiocapsouza@gmail.com>
Date: Sun, 15 Nov 2020 12:30:46 -0300
Subject: [PATCH] Revert "Clean Trainer tests and datasets dep (#8268)"

This reverts commit c03c385a9424018ba20360623003bdf060c66523.
---
 .circleci/config.yml  |  6 ++++++
 tests/test_trainer.py | 31 +++++++++++++------------------
 2 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 662e18c6953dae..524e5a940d7053 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -77,6 +77,7 @@ jobs:
                       - v0.4-torch_and_tf-{{ checksum "setup.py" }}
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
+            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,tf-cpu,torch,testing]
             - save_cache:
                 key: v0.4-{{ checksum "setup.py" }}
@@ -101,6 +102,7 @@ jobs:
                       - v0.4-torch-{{ checksum "setup.py" }}
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
+            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,torch,testing]
             - save_cache:
                   key: v0.4-torch-{{ checksum "setup.py" }}
@@ -127,6 +129,7 @@ jobs:
                       - v0.4-tf-{{ checksum "setup.py" }}
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
+            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,tf-cpu,testing]
             - save_cache:
                   key: v0.4-tf-{{ checksum "setup.py" }}
@@ -151,6 +154,7 @@ jobs:
                     - v0.4-flax-{{ checksum "setup.py" }}
                     - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
+            - run: pip install git+https://github.com/huggingface/datasets
             - run: sudo pip install .[flax,sklearn,torch,testing]
             - save_cache:
                   key: v0.4-flax-{{ checksum "setup.py" }}
@@ -175,6 +179,7 @@ jobs:
                       - v0.4-torch-{{ checksum "setup.py" }}
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
+            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,torch,testing]
             - save_cache:
                   key: v0.4-torch-{{ checksum "setup.py" }}
@@ -199,6 +204,7 @@ jobs:
                       - v0.4-tf-{{ checksum "setup.py" }}
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
+            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,tf-cpu,testing]
             - save_cache:
                   key: v0.4-tf-{{ checksum "setup.py" }}
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
index 497c6c8b4daa73..6505539cdac042 100644
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -18,13 +18,13 @@
 import tempfile
 import unittest
 
+import datasets
 import numpy as np
 
 from transformers import AutoTokenizer, EvaluationStrategy, PretrainedConfig, TrainingArguments, is_torch_available
 from transformers.file_utils import WEIGHTS_NAME
 from transformers.testing_utils import (
     get_tests_dir,
-    require_datasets,
     require_optuna,
     require_sentencepiece,
     require_tokenizers,
@@ -340,10 +340,7 @@ def test_predict(self):
         self.assertTrue(np.array_equal(labels[0], trainer.eval_dataset.ys[0]))
         self.assertTrue(np.array_equal(labels[1], trainer.eval_dataset.ys[1]))
 
-    @require_datasets
     def test_trainer_with_datasets(self):
-        import datasets
-
         np.random.seed(42)
         x = np.random.normal(size=(64,)).astype(np.float32)
         y = 2.0 * x + 3.0 + np.random.normal(scale=0.1, size=(64,))
@@ -661,17 +658,15 @@ def model_init(trial):
         def hp_name(trial):
             return MyTrialShortNamer.shortname(trial.params)
 
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            trainer = get_regression_trainer(
-                output_dir=tmp_dir,
-                learning_rate=0.1,
-                logging_steps=1,
-                evaluation_strategy=EvaluationStrategy.EPOCH,
-                num_train_epochs=4,
-                disable_tqdm=True,
-                load_best_model_at_end=True,
-                logging_dir="runs",
-                run_name="test",
-                model_init=model_init,
-            )
-            trainer.hyperparameter_search(direction="minimize", hp_space=hp_space, hp_name=hp_name, n_trials=4)
+        trainer = get_regression_trainer(
+            learning_rate=0.1,
+            logging_steps=1,
+            evaluation_strategy=EvaluationStrategy.EPOCH,
+            num_train_epochs=4,
+            disable_tqdm=True,
+            load_best_model_at_end=True,
+            logging_dir="runs",
+            run_name="test",
+            model_init=model_init,
+        )
+        trainer.hyperparameter_search(direction="minimize", hp_space=hp_space, hp_name=hp_name, n_trials=4)