From bd3d50d81d43d5815ab72d0d1525a7cf6a3bed6e Mon Sep 17 00:00:00 2001
From: Sean Naren <sean.narenthiran@gmail.com>
Date: Mon, 1 Feb 2021 18:00:19 +0000
Subject: [PATCH] Update API to work with latest transformers/datasets package
 (#40)

* Update API to work with latest transformers/datasets package

* Move logic to textx classification pipeline

* Strong typing
---
 flash/core/classification.py       |  5 +++--
 flash/text/classification/data.py  | 10 ++++++++--
 flash/text/classification/model.py |  3 ++-
 requirements.txt                   |  4 ++--
 4 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/flash/core/classification.py b/flash/core/classification.py
index 0800a534bf..e71f1c74d0 100644
--- a/flash/core/classification.py
+++ b/flash/core/classification.py
@@ -11,9 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any
+from typing import Any, Union
 
 import torch
+from transformers.modeling_outputs import SequenceClassifierOutput
 
 from flash.core.data import TaskDataPipeline
 from flash.core.model import Task
@@ -21,7 +22,7 @@
 
 class ClassificationDataPipeline(TaskDataPipeline):
 
-    def before_uncollate(self, batch: torch.Tensor) -> torch.Tensor:
+    def before_uncollate(self, batch: Union[torch.Tensor, tuple]) -> torch.Tensor:
         if isinstance(batch, tuple):
             batch = batch[0]
         return torch.softmax(batch, -1)
diff --git a/flash/text/classification/data.py b/flash/text/classification/data.py
index a491cb3994..4ae0f7e768 100644
--- a/flash/text/classification/data.py
+++ b/flash/text/classification/data.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from functools import partial
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, Union
 
 import torch
 from datasets import load_dataset
@@ -20,6 +20,7 @@
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from torch import Tensor
 from transformers import AutoTokenizer, default_data_collator
+from transformers.modeling_outputs import SequenceClassifierOutput
 
 from flash.core.classification import ClassificationDataPipeline
 from flash.core.data import DataModule
@@ -48,7 +49,6 @@ def prepare_dataset(
     label_to_class_mapping=None,
     predict=False,
 ):
-
     data_files = {}
 
     if train_file is not None:
@@ -141,6 +141,12 @@ def after_collate(self, batch: Tensor) -> Tensor:
                 batch["input_ids"] = batch["input_ids"].squeeze(0)
         return batch
 
+    def before_uncollate(self, batch: Union[torch.Tensor, tuple,
+                                            SequenceClassifierOutput]) -> Union[tuple, torch.Tensor]:
+        if isinstance(batch, SequenceClassifierOutput):
+            batch = batch.logits
+        return super().before_uncollate(batch)
+
 
 class TextClassificationData(DataModule):
     """Data module for text classification tasks."""
diff --git a/flash/text/classification/model.py b/flash/text/classification/model.py
index 542b82f2fe..eff2bfa050 100644
--- a/flash/text/classification/model.py
+++ b/flash/text/classification/model.py
@@ -69,7 +69,8 @@ def forward(self, batch_dict):
 
     def step(self, batch, batch_idx) -> dict:
         output = {}
-        loss, logits = self.forward(batch)
+        out = self.forward(batch)
+        loss, logits = out[:2]
         output["loss"] = loss
         output["y_hat"] = logits
         probs = self.data_pipeline.before_uncollate(logits)
diff --git a/requirements.txt b/requirements.txt
index bf6e8dc08f..9182f75e0c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,8 +3,8 @@ torch==1.7.1
 PyYAML==5.3.1
 Pillow>=7.2
 torchvision==0.8.2
-transformers==3.1.0
-datasets==1.0.1
+transformers==4.2.2
+datasets==1.2.1
 pandas==1.1.2
 scikit-learn==0.24.0
 numpy