Repository: flairNLP/flair

Issue:
loss_weight parameter causes exception for TextPairClassifier
**Describe the bug**
Using loss_weight parameter for TextPairClassier causes the trainer to fail when it is doing a final evaluation at the end of training

**To Reproduce**
```
tagger = TextPairClassifier(
    document_embeddings=embedding,
    label_dictionary=label_dictionary,
    label_type='similar', loss_weights = {'1': 5., '0': 0.25 }
)


trainer = ModelTrainer(tagger, corpus)

trainer.train('./mod',
              learning_rate=2e-4,
              mini_batch_chunk_size=16, # this can be removed if you hae a big GPU
              train_with_dev=False,
              max_epochs=1, embeddings_storage_mode='cpu', checkpoint=True, num_workers=12, scheduler=OneCycleLR, optimizer=torch.optim.AdamW)
```


**Expected behavior**
The model should train without exception

**Screenshots**
If applicable, add screenshots to help explain your problem.

**Environment (please complete the following information):**
 - OS [e.g. iOS, Linux]:
 - Version [e.g. flair-0.3.2]:

**Additional context**
stacktrace
```
AttributeError                            Traceback (most recent call last)
<ipython-input-9-59bda6c3483a> in <module>
      9               mini_batch_chunk_size=16, # this can be removed if you hae a big GPU
     10               train_with_dev=False,
---> 11               max_epochs=1, embeddings_storage_mode='cpu', checkpoint=True, num_workers=12, scheduler=OneCycleLR, optimizer=torch.optim.AdamW)

~\.conda\envs\knlu\lib\site-packages\flair\trainers\trainer.py in train(self, base_path, learning_rate, mini_batch_size, mini_batch_chunk_size, max_epochs, train_with_dev, train_with_test, monitor_train, monitor_test, main_evaluation_metric, scheduler, anneal_factor, patience, min_learning_rate, initial_extra_patience, optimizer, cycle_momentum, warmup_fraction, embeddings_storage_mode, checkpoint, save_final_model, anneal_with_restarts, anneal_with_prestarts, anneal_against_dev_loss, batch_growth_annealing, shuffle, param_selection_mode, write_weights, num_workers, sampler, use_amp, amp_opt_level, eval_on_train_fraction, eval_on_train_shuffle, save_model_each_k_epochs, tensorboard_comment, use_swa, use_final_model_for_eval, gold_label_dictionary_for_eval, create_file_logs, create_loss_file, epoch, use_tensorboard, tensorboard_log_dir, metrics_for_tensorboard, optimizer_state_dict, scheduler_state_dict, save_optimizer_state, **kwargs)
    748                 num_workers=num_workers,
    749                 main_evaluation_metric=main_evaluation_metric,
--> 750                 gold_label_dictionary_for_eval=gold_label_dictionary_for_eval,
    751             )
    752         else:

~\.conda\envs\knlu\lib\site-packages\flair\trainers\trainer.py in final_test(self, base_path, eval_mini_batch_size, main_evaluation_metric, num_workers, gold_label_dictionary_for_eval)
    836 
    837         if (base_path / "best-model.pt").exists():
--> 838             self.model.load_state_dict(self.model.load(base_path / "best-model.pt").state_dict())
    839         else:
    840             log.info("Testing using last state of model ...")

~\.conda\envs\knlu\lib\site-packages\flair\nn\model.py in load(cls, model)
    133             state = torch.load(f, map_location='cpu')
    134 
--> 135         model = cls._init_model_with_state_dict(state)
    136 
    137         if 'model_card' in state:

~\.conda\envs\knlu\lib\site-packages\flair\models\pairwise_classification_model.py in _init_model_with_state_dict(state)
    149             multi_label_threshold=0.5 if "multi_label_threshold" not in state.keys() else state["multi_label_threshold"],
    150             loss_weights=state["loss_weights"],
--> 151             embed_separately=state["embed_separately"],
    152         )
    153         model.load_state_dict(state["state_dict"])

~\.conda\envs\knlu\lib\site-packages\flair\models\pairwise_classification_model.py in __init__(self, document_embeddings, label_type, embed_separately, **classifierargs)
     34         (if any label's weight is unspecified it will default to 1.0)
     35         """
---> 36         super().__init__(**classifierargs)
     37 
     38         self.document_embeddings: flair.embeddings.DocumentEmbeddings = document_embeddings

~\.conda\envs\knlu\lib\site-packages\flair\nn\model.py in __init__(self, label_dictionary, multi_label, multi_label_threshold, loss_weights)
    461             weight_list = [1.0 for i in range(n_classes)]
    462             for i, tag in enumerate(self.label_dictionary.get_items()):
--> 463                 if tag in loss_weights.keys():
    464                     weight_list[i] = loss_weights[tag]
    465             self.loss_weights = torch.FloatTensor(weight_list).to(flair.device)

AttributeError: 'Tensor' object has no attribute 'keys'
```

loss_weight parameter causes exception for TextPairClassifier
**Describe the bug**
Using loss_weight parameter for TextPairClassier causes the trainer to fail when it is doing a final evaluation at the end of training

**To Reproduce**
```
tagger = TextPairClassifier(
    document_embeddings=embedding,
    label_dictionary=label_dictionary,
    label_type='similar', loss_weights = {'1': 5., '0': 0.25 }
)


trainer = ModelTrainer(tagger, corpus)

trainer.train('./mod',
              learning_rate=2e-4,
              mini_batch_chunk_size=16, # this can be removed if you hae a big GPU
              train_with_dev=False,
              max_epochs=1, embeddings_storage_mode='cpu', checkpoint=True, num_workers=12, scheduler=OneCycleLR, optimizer=torch.optim.AdamW)
```


**Expected behavior**
The model should train without exception

**Screenshots**
If applicable, add screenshots to help explain your problem.

**Environment (please complete the following information):**
 - OS [e.g. iOS, Linux]:
 - Version [e.g. flair-0.3.2]:

**Additional context**
stacktrace
```
AttributeError                            Traceback (most recent call last)
<ipython-input-9-59bda6c3483a> in <module>
      9               mini_batch_chunk_size=16, # this can be removed if you hae a big GPU
     10               train_with_dev=False,
---> 11               max_epochs=1, embeddings_storage_mode='cpu', checkpoint=True, num_workers=12, scheduler=OneCycleLR, optimizer=torch.optim.AdamW)

~\.conda\envs\knlu\lib\site-packages\flair\trainers\trainer.py in train(self, base_path, learning_rate, mini_batch_size, mini_batch_chunk_size, max_epochs, train_with_dev, train_with_test, monitor_train, monitor_test, main_evaluation_metric, scheduler, anneal_factor, patience, min_learning_rate, initial_extra_patience, optimizer, cycle_momentum, warmup_fraction, embeddings_storage_mode, checkpoint, save_final_model, anneal_with_restarts, anneal_with_prestarts, anneal_against_dev_loss, batch_growth_annealing, shuffle, param_selection_mode, write_weights, num_workers, sampler, use_amp, amp_opt_level, eval_on_train_fraction, eval_on_train_shuffle, save_model_each_k_epochs, tensorboard_comment, use_swa, use_final_model_for_eval, gold_label_dictionary_for_eval, create_file_logs, create_loss_file, epoch, use_tensorboard, tensorboard_log_dir, metrics_for_tensorboard, optimizer_state_dict, scheduler_state_dict, save_optimizer_state, **kwargs)
    748                 num_workers=num_workers,
    749                 main_evaluation_metric=main_evaluation_metric,
--> 750                 gold_label_dictionary_for_eval=gold_label_dictionary_for_eval,
    751             )
    752         else:

~\.conda\envs\knlu\lib\site-packages\flair\trainers\trainer.py in final_test(self, base_path, eval_mini_batch_size, main_evaluation_metric, num_workers, gold_label_dictionary_for_eval)
    836 
    837         if (base_path / "best-model.pt").exists():
--> 838             self.model.load_state_dict(self.model.load(base_path / "best-model.pt").state_dict())
    839         else:
    840             log.info("Testing using last state of model ...")

~\.conda\envs\knlu\lib\site-packages\flair\nn\model.py in load(cls, model)
    133             state = torch.load(f, map_location='cpu')
    134 
--> 135         model = cls._init_model_with_state_dict(state)
    136 
    137         if 'model_card' in state:

~\.conda\envs\knlu\lib\site-packages\flair\models\pairwise_classification_model.py in _init_model_with_state_dict(state)
    149             multi_label_threshold=0.5 if "multi_label_threshold" not in state.keys() else state["multi_label_threshold"],
    150             loss_weights=state["loss_weights"],
--> 151             embed_separately=state["embed_separately"],
    152         )
    153         model.load_state_dict(state["state_dict"])

~\.conda\envs\knlu\lib\site-packages\flair\models\pairwise_classification_model.py in __init__(self, document_embeddings, label_type, embed_separately, **classifierargs)
     34         (if any label's weight is unspecified it will default to 1.0)
     35         """
---> 36         super().__init__(**classifierargs)
     37 
     38         self.document_embeddings: flair.embeddings.DocumentEmbeddings = document_embeddings

~\.conda\envs\knlu\lib\site-packages\flair\nn\model.py in __init__(self, label_dictionary, multi_label, multi_label_threshold, loss_weights)
    461             weight_list = [1.0 for i in range(n_classes)]
    462             for i, tag in enumerate(self.label_dictionary.get_items()):
--> 463                 if tag in loss_weights.keys():
    464                     weight_list[i] = loss_weights[tag]
    465             self.loss_weights = torch.FloatTensor(weight_list).to(flair.device)

AttributeError: 'Tensor' object has no attribute 'keys'
```



Discussion:
To solve this issue at my end, I made my own TextPairClassifier class with the following methods. 

```
def _get_state_dict(self):
        model_state = {
            "state_dict": self.state_dict(),
            "document_embeddings": self.document_embeddings,
            "label_dictionary": self.label_dictionary,
            "label_type": self.label_type,
            "multi_label": self.multi_label,
            "multi_label_threshold": self.multi_label_threshold,
            "weight_dict": self.weight_dict,
            "embed_separately": self.embed_separately,
        }
        return model_state
```

and

```
 @staticmethod
    def _init_model_with_state_dict(state):
       weights = None if "weight_dict" not in state.keys() else state["weight_dict"]
        model = TextPairClassifier(
            document_embeddings=state["document_embeddings"],
            label_dictionary=state["label_dictionary"],
            label_type=state["label_type"],
            multi_label=state["multi_label"],
            multi_label_threshold=0.5 if "multi_label_threshold" not in state.keys() else state["multi_label_threshold"],
            loss_weights=weights,
            embed_separately=state["embed_separately"],
        )
        model.load_state_dict(state["state_dict"])
        return model
```

@lukasgarbas can you check this bug?
I confirm the bug with `RelationExtractor`. It can be fixed in `DefaultClassifier` `init` by checking if weights are a dict or a tensor.
```
        if isinstance(loss_weights, dict):
            n_classes = len(self.label_dictionary)
            weight_list = [1.0 for i in range(n_classes)]
            for i, tag in enumerate(self.label_dictionary.get_items()):
                if tag in loss_weights.keys():
                    weight_list[i] = loss_weights[tag]
            self.loss_weights = torch.FloatTensor(weight_list).to(flair.device)
        elif isinstance(loss_weights, torch.FloatTensor):
            self.loss_weights = loss_weights
        else:
            self.loss_weights = None
```
To solve this issue at my end, I made my own TextPairClassifier class with the following methods. 

```
def _get_state_dict(self):
        model_state = {
            "state_dict": self.state_dict(),
            "document_embeddings": self.document_embeddings,
            "label_dictionary": self.label_dictionary,
            "label_type": self.label_type,
            "multi_label": self.multi_label,
            "multi_label_threshold": self.multi_label_threshold,
            "weight_dict": self.weight_dict,
            "embed_separately": self.embed_separately,
        }
        return model_state
```

and

```
 @staticmethod
    def _init_model_with_state_dict(state):
       weights = None if "weight_dict" not in state.keys() else state["weight_dict"]
        model = TextPairClassifier(
            document_embeddings=state["document_embeddings"],
            label_dictionary=state["label_dictionary"],
            label_type=state["label_type"],
            multi_label=state["multi_label"],
            multi_label_threshold=0.5 if "multi_label_threshold" not in state.keys() else state["multi_label_threshold"],
            loss_weights=weights,
            embed_separately=state["embed_separately"],
        )
        model.load_state_dict(state["state_dict"])
        return model
```

@lukasgarbas can you check this bug?
I confirm the bug with `RelationExtractor`. It can be fixed in `DefaultClassifier` `init` by checking if weights are a dict or a tensor.
```
        if isinstance(loss_weights, dict):
            n_classes = len(self.label_dictionary)
            weight_list = [1.0 for i in range(n_classes)]
            for i, tag in enumerate(self.label_dictionary.get_items()):
                if tag in loss_weights.keys():
                    weight_list[i] = loss_weights[tag]
            self.loss_weights = torch.FloatTensor(weight_list).to(flair.device)
        elif isinstance(loss_weights, torch.FloatTensor):
            self.loss_weights = loss_weights
        else:
            self.loss_weights = None
```

Resolving patch:
diff --git a/flair/datasets/text_text.py b/flair/datasets/text_text.py
--- a/flair/datasets/text_text.py
+++ b/flair/datasets/text_text.py
@@ -502,6 +502,7 @@ def __init__(
 
         self.eval_dataset = DataPairDataset(
             data_folder / "RTE/eval_dataset.tsv",
+            label_type=label_type,
             columns=[1, 2, 3],
             use_tokenizer=use_tokenizer,
             max_tokens_per_doc=max_tokens_per_doc,
@@ -621,6 +622,7 @@ def __init__(
 
         self.eval_dataset = DataPairDataset(
             data_folder / "MNLI" / eval_dataset,
+            label_type=label_type,
             columns=[8, 9, 11],
             use_tokenizer=use_tokenizer,
             max_tokens_per_doc=max_tokens_per_doc,
@@ -791,6 +793,7 @@ def __init__(
 
         self.eval_dataset = DataPairDataset(
             data_folder / "QNLI/eval_dataset.tsv",
+            label_type=label_type,
             columns=[1, 2, 3],
             use_tokenizer=use_tokenizer,
             max_tokens_per_doc=max_tokens_per_doc,
@@ -878,6 +881,7 @@ def __init__(
 
         self.eval_dataset = DataPairDataset(
             data_folder / "QQP/eval_dataset.tsv",
+            label_type=label_type,
             columns=[1, 2, 0],
             use_tokenizer=use_tokenizer,
             max_tokens_per_doc=max_tokens_per_doc,
@@ -965,6 +969,7 @@ def __init__(
 
         self.eval_dataset = DataPairDataset(
             data_folder / "WNLI/eval_dataset.tsv",
+            label_type=label_type,
             columns=[1, 2, 3],
             use_tokenizer=use_tokenizer,
             max_tokens_per_doc=max_tokens_per_doc,
@@ -1053,6 +1058,7 @@ def __init__(
 
         self.eval_dataset = DataPairDataset(
             data_folder / "RTE/eval_dataset.tsv",
+            label_type="textual_entailment",
             columns=[0, 1, 2],
             use_tokenizer=use_tokenizer,
             max_tokens_per_doc=max_tokens_per_doc,
diff --git a/flair/models/pairwise_classification_model.py b/flair/models/pairwise_classification_model.py
--- a/flair/models/pairwise_classification_model.py
+++ b/flair/models/pairwise_classification_model.py
@@ -142,7 +142,7 @@ def _get_state_dict(self):
             "label_type": self.label_type,
             "multi_label": self.multi_label,
             "multi_label_threshold": self.multi_label_threshold,
-            "loss_weights": self.loss_weights,
+            "weight_dict": self.weight_dict,
             "embed_separately": self.embed_separately,
         }
         return model_state
@@ -158,7 +158,7 @@ def _init_model_with_state_dict(state):
             multi_label_threshold=0.5
             if "multi_label_threshold" not in state.keys()
             else state["multi_label_threshold"],
-            loss_weights=state["loss_weights"],
+            loss_weights=state["weight_dict"],
             embed_separately=state["embed_separately"],
         )
         model.load_state_dict(state["state_dict"])
diff --git a/flair/models/relation_extractor_model.py b/flair/models/relation_extractor_model.py
--- a/flair/models/relation_extractor_model.py
+++ b/flair/models/relation_extractor_model.py
@@ -259,7 +259,7 @@ def _get_state_dict(self):
             "label_dictionary": self.label_dictionary,
             "label_type": self.label_type,
             "entity_label_type": self.entity_label_type,
-            "loss_weights": self.loss_weights,
+            "weight_dict": self.weight_dict,
             "pooling_operation": self.pooling_operation,
             "dropout_value": self.dropout_value,
             "locked_dropout_value": self.locked_dropout_value,
@@ -276,7 +276,7 @@ def _init_model_with_state_dict(state):
             label_dictionary=state["label_dictionary"],
             label_type=state["label_type"],
             entity_label_type=state["entity_label_type"],
-            loss_weights=state["loss_weights"],
+            loss_weights=state["weight_dict"],
             pooling_operation=state["pooling_operation"],
             dropout_value=state["dropout_value"],
             locked_dropout_value=state["locked_dropout_value"],


Test patch: