update setup.py

autox/autox_nlp/feature_engineer/nlp_feature.py

@@ -8,12 +8,16 @@ from sklearn.model_selection import KFold
 from sklearn.cluster import KMeans
 from scipy import sparse
 from sklearn.feature_extraction.text import TfidfVectorizer
+from torch.utils.data import DataLoader, Dataset
+try:
+    from datasets import Dataset as DataSet
+except Exception as e:
+    os.system('pip install datasets')
+    from datasets import Dataset as DataSet
 from transformers import (AutoModel, AutoModelForMaskedLM,
                           AutoTokenizer, LineByLineTextDataset,
                           DataCollatorForLanguageModeling,
                           Trainer, TrainingArguments, PreTrainedTokenizerFast, pipeline)
-from torch.utils.data import DataLoader, Dataset
-from datasets import Dataset as DataSet
 from sklearn.linear_model import Ridge
 from tokenizers import (
     decoders,

setup.py

@@ -23,10 +23,9 @@ setup(name="automl-x",
             'keras',
             'tensorflow',
             'gensim',
-            'pyarrow>=6.0.0'
             # 'glove-python-binary',
             'transformers',
-            'datasets',
+            # 'datasets',
             'pyarrow>=6.0.0',
         ],
         python_requires='>=3.6',