Repository: pytorch/ignite

Issue:
Restrict model options in transformers examples
In transformers example https://github.com/pytorch/ignite/tree/master/examples/contrib/transformers it's up to user to override the default model which is `bert-base-uncased`, a lot of models take similar inputs to `BERT` and similar outputs too, but models like `distilbert-base-uncased`, `distilroberta-base`, `bart-base` (and many other models) will not work here as they work on a bit different way regarding to the inputs and outputs of the model, check here for more info: https://huggingface.co/transformers/model_doc/distilbert.html#distilbertmodel

So we will get an error similar to this while using `DistilBERT`
```
KeyError: 'token_type_ids'
```
And also if we used something like `XLNet` that won't work well and we will have dimensions issue, because `XLNet` doesn't return a `pooler_output`, check here https://huggingface.co/transformers/model_doc/xlnet.html#transformers.XLNetModel

So what we should do is:
- We don't provide the option of choosing models, and we just keep BERT, and if the user wants to change the model, he can do that manually
- Keep it that way and mention in the docs that, the models that can be used here are `BERT`, `RoBERTa` (and if there is others we mentions them).

EDIT: After deciding to go with the first option, what we need to do now is to remove `model` from the argument, and set `config['model'] = 'bert-base-uncased'` manually inside `run` method, and mention in the docs we are using `BERT` by default, and maybe leave a note about if the user wants to try another model, he should that himself.


Discussion:
Thanks for reporting that @KickItLikeShika , I think the first option is good enough.

Could you please provide few links on the code and explain in details on how to do that. Such that someone new could try to tackle the issue. Thanks
@vfdev-5 Updated the description.
@KickItLikeShika Another way I think we can fix this is by passing the dictionary that tokenizer returns instead of making our own dictionary in `dataset.py`. So models which don't require `token_type_ids` their tokenizer will not produce them also.  Have a look at [this](https://github.com/huggingface/transformers/blob/master/examples/pytorch/text-classification/run_glue_no_trainer.py#L297-L311) example of Classification from the HF repo.
@Ishan-Kumar2 seems like a good idea for resolving the input issues, but we still have the same issue related to the Transformer output, most of other transformers don't have a `pooled_output`, I don't see how your suggestion can solve this, what do you think?
We can replace the model with `AutoModelForSequenceClassification` and then have the loss computed like this 
```python
outputs = model(**inputs, labels=labels)
loss = outputs.loss
logits = output.logits
```
This should work for both, looking at these examples for [XLNet](https://huggingface.co/transformers/model_doc/xlnet.html#transformers.XLNetForSequenceClassification) and [Bert](https://huggingface.co/transformers/model_doc/bert.html#transformers.BertForSequenceClassification) right?
Initially it looks good, what do you think @vfdev-5? if that's OK, can you try this out please? @Ishan-Kumar2 
OK for me if you think it could be understandable by other users.

Resolving patch:
diff --git a/examples/contrib/transformers/dataset.py b/examples/contrib/transformers/dataset.py
--- a/examples/contrib/transformers/dataset.py
+++ b/examples/contrib/transformers/dataset.py
@@ -11,24 +11,19 @@ def __init__(self, texts, labels, tokenizer, max_length):
     def __getitem__(self, idx):
         text = str(self.texts[idx])
         text = " ".join(text.split())
-        inputs = self.tokenizer.encode_plus(
-            text, None, add_special_tokens=True, max_length=self.max_length, truncation=True
+        inputs = self.tokenizer(
+            text,
+            None,
+            add_special_tokens=True,
+            max_length=self.max_length,
+            truncation=True,
+            padding="max_length",
+            return_tensors="pt",
         )
+        inputs = {k: v.type(torch.long).squeeze(0) for k, v in inputs.items()}
 
-        ids = inputs["input_ids"]
-        token_type_ids = inputs["token_type_ids"]
-        mask = inputs["attention_mask"]
-        padding_length = self.max_length - len(ids)
-
-        ids = ids + ([0] * padding_length)
-        mask = mask + ([0] * padding_length)
-        token_type_ids = token_type_ids + ([0] * padding_length)
-        return {
-            "input_ids": torch.tensor(ids, dtype=torch.long),
-            "attention_mask": torch.tensor(mask, dtype=torch.long),
-            "token_type_ids": torch.tensor(token_type_ids, dtype=torch.long),
-            "label": torch.tensor(self.labels[idx], dtype=torch.float),
-        }
+        labels_pt = torch.tensor(self.labels[idx], dtype=torch.float)
+        return inputs, labels_pt
 
     def __len__(self):
         return len(self.labels)
diff --git a/examples/contrib/transformers/main.py b/examples/contrib/transformers/main.py
--- a/examples/contrib/transformers/main.py
+++ b/examples/contrib/transformers/main.py
@@ -144,7 +144,7 @@ def run(
     dropout=0.3,
     n_fc=768,
     max_length=256,
-    batch_size=128,
+    batch_size=32,
     weight_decay=0.01,
     num_workers=4,
     num_epochs=3,
@@ -303,21 +303,17 @@ def create_trainer(model, optimizer, criterion, lr_scheduler, train_sampler, con
 
     def train_step(engine, batch):
 
-        input_ids = batch["input_ids"]
-        attention_mask = batch["attention_mask"]
-        token_type_ids = batch["token_type_ids"]
-        labels = batch["label"].view(-1, 1)
+        input_batch = batch[0]
+        labels = batch[1].view(-1, 1)
 
-        if input_ids.device != device:
-            input_ids = input_ids.to(device, non_blocking=True, dtype=torch.long)
-            attention_mask = attention_mask.to(device, non_blocking=True, dtype=torch.long)
-            token_type_ids = token_type_ids.to(device, non_blocking=True, dtype=torch.long)
+        if labels.device != device:
+            input_batch = {k: v.to(device, non_blocking=True, dtype=torch.long) for k, v in batch[0].items()}
             labels = labels.to(device, non_blocking=True, dtype=torch.float)
 
         model.train()
 
         with autocast(enabled=with_amp):
-            y_pred = model(input_ids, attention_mask, token_type_ids)
+            y_pred = model(input_batch)
             loss = criterion(y_pred, labels)
 
         optimizer.zero_grad()
@@ -372,19 +368,16 @@ def create_evaluator(model, metrics, config, tag="val"):
     @torch.no_grad()
     def evaluate_step(engine, batch):
         model.eval()
-        input_ids = batch["input_ids"]
-        attention_mask = batch["attention_mask"]
-        token_type_ids = batch["token_type_ids"]
-        labels = batch["label"].view(-1, 1)
-
-        if input_ids.device != device:
-            input_ids = input_ids.to(device, non_blocking=True, dtype=torch.long)
-            attention_mask = attention_mask.to(device, non_blocking=True, dtype=torch.long)
-            token_type_ids = token_type_ids.to(device, non_blocking=True, dtype=torch.long)
+
+        input_batch = batch[0]
+        labels = batch[1].view(-1, 1)
+
+        if labels.device != device:
+            input_batch = {k: v.to(device, non_blocking=True, dtype=torch.long) for k, v in batch[0].items()}
             labels = labels.to(device, non_blocking=True, dtype=torch.float)
 
         with autocast(enabled=with_amp):
-            output = model(input_ids, attention_mask, token_type_ids)
+            output = model(input_batch)
         return output, labels
 
     evaluator = Engine(evaluate_step)
diff --git a/examples/contrib/transformers/model.py b/examples/contrib/transformers/model.py
--- a/examples/contrib/transformers/model.py
+++ b/examples/contrib/transformers/model.py
@@ -1,19 +1,22 @@
 import torch.nn as nn
-from transformers import AutoModel
+from transformers import AutoConfig, AutoModelForSequenceClassification
 
 
 class TransformerModel(nn.Module):
     def __init__(self, model_name, model_dir, dropout, n_fc, n_classes):
         super(TransformerModel, self).__init__()
-        self.transformer = AutoModel.from_pretrained(model_name, cache_dir=model_dir)
-        self.drop = nn.Dropout(dropout)
-        self.classifier = nn.Linear(n_fc, n_classes)
+        self.config = AutoConfig.from_pretrained(
+            model_name,
+            num_labels=n_classes,
+            output_hidden_states=n_fc,
+            classifier_dropout=dropout,
+            output_attentions=True,
+        )
+        self.transformer = AutoModelForSequenceClassification.from_pretrained(
+            model_name, cache_dir=model_dir, config=self.config
+        )
 
-    def forward(self, ids, mask, token_type_ids):
+    def forward(self, inputs):
+        output = self.transformer(**inputs)["logits"]
 
-        hidden_output, pooled_output = self.transformer(
-            ids, attention_mask=mask, token_type_ids=token_type_ids, return_dict=False
-        )
-        pooled_output = self.drop(pooled_output)
-        output = self.classifier(pooled_output)
         return output


Test patch: