Fix return all score issue (#70)

* Use self.tokenizer_name in text_classification

* Fix distilbert classification input for return_all_scores

* Fix post_process output for return_all_scores in text_classification.py

* Fix roberta input for return_all_scores in roberta_text_classification.py

* Fix bert input for return_all_scores in bert_text_classification.py

hezar/models/text_classification/bert/bert_text_classification.py

@@ -42,19 +42,32 @@ def _build_inner_config(self):
     def forward(self, inputs, **kwargs) -> Dict:
         input_ids = inputs.get("token_ids")
         attention_mask = inputs.get("attention_mask", None)
+        token_types_ids = inputs.get("token_type_ids", None) or kwargs.get("token_type_ids", None)
+        position_ids = inputs.get("position_ids", None) or kwargs.get("position_ids", None)
         head_mask = inputs.get("head_mask", None)
         inputs_embeds = inputs.get("inputs_embeds", None)
+        encoder_hidden_states = inputs.get("encoder_hidden_states", None) or kwargs.get("encoder_hidden_states", None)
+        encoder_attention_mask = inputs.get("encoder_attention_mask", None) or kwargs.get("encoder_attention_mask", None)
+        past_key_values = inputs.get("past_key_values", None) or kwargs.get("past_key_values", None)
+        use_cache = inputs.get("use_cache", None) or kwargs.get("use_cache", None)
         output_attentions = inputs.get("output_attentions", None)
         output_hidden_states = inputs.get("output_hidden_states", None)
+        return_dict = inputs.get("return_dict", None) or kwargs.get("return_dict", None)
 
         lm_outputs = self.bert(
             input_ids,
             attention_mask=attention_mask,
+            token_type_ids=token_types_ids,
+            position_ids=position_ids,
             head_mask=head_mask,
             inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
-            **kwargs,
+            return_dict=return_dict,
         )
         pooled_output = lm_outputs[1]
         pooled_output = self.dropout(pooled_output)

hezar/models/text_classification/distilbert/distilbert_text_classification.py

@@ -43,6 +43,7 @@ def forward(self, inputs, **kwargs) -> Dict:
         inputs_embeds = inputs.get("inputs_embeds", None)
         output_attentions = inputs.get("output_attentions", None)
         output_hidden_states = inputs.get("output_hidden_states", None)
+        return_dict = inputs.get("return_dict", None) or kwargs.get("return_dict", None)
 
         lm_outputs = self.distilbert(
             input_ids,
@@ -51,7 +52,7 @@ def forward(self, inputs, **kwargs) -> Dict:
             inputs_embeds=inputs_embeds,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
-            **kwargs,
+            return_dict=return_dict,
         )
         hidden_state = lm_outputs[0]
         pooled_output = hidden_state[:, 0]

hezar/models/text_classification/roberta/roberta_text_classification.py

@@ -36,19 +36,32 @@ def _build_inner_config(self):
     def forward(self, inputs, **kwargs):
         input_ids = inputs.get("token_ids")
         attention_mask = inputs.get("attention_mask", None)
+        token_type_ids = inputs.get("token_type_ids", None) or kwargs.get("token_type_ids", None)
+        position_ids = inputs.get("position_ids", None) or kwargs.get("position_ids", None)
         head_mask = inputs.get("head_mask", None)
         inputs_embeds = inputs.get("inputs_embeds", None)
+        encoder_hidden_states = inputs.get("encoder_hidden_states", None) or kwargs.get("encoder_hidden_states", None)
+        encoder_attention_mask = inputs.get("encoder_attention_mask", None) or kwargs.get("encoder_attention_mask", None)
+        past_key_values = inputs.get("past_key_values", None) or kwargs.get("past_key_values", None)
+        use_cache = inputs.get("use_cache", None) or kwargs.get("use_cache", None)
         output_attentions = inputs.get("output_attentions", None)
         output_hidden_states = inputs.get("output_hidden_states", None)
+        return_dict = inputs.get("return_dict", None) or kwargs.get("return_dict", None)
 
         lm_outputs = self.roberta(
             input_ids,
             attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
             head_mask=head_mask,
             inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
-            **kwargs,
+            return_dict=return_dict,
         )
         sequence_output = lm_outputs[0]
         logits = self.classifier(sequence_output)

hezar/models/text_classification/text_classification.py

@@ -14,7 +14,7 @@ def preprocess(self, inputs: Union[str, List[str]], **kwargs):
         if "text_normalizer" in self.preprocessor:
             normalizer = self.preprocessor["text_normalizer"]
             inputs = normalizer(inputs)
-        tokenizer = self.preprocessor[self._tokenizer_name]
+        tokenizer = self.preprocessor[self.tokenizer_name]
         inputs = tokenizer(inputs, return_tensors="pt", device=self.device)
         return inputs
 
@@ -27,9 +27,9 @@ def post_process(self, inputs, **kwargs) -> Dict:
             outputs = []
             for sample_index in range(predictions.shape[0]):
                 sample_outputs = []
-                for prediction, prob in zip(predictions[sample_index], predictions_probs[sample_index]):
-                    label = self.config.id2label[prediction.item()]
-                    sample_outputs.append({"label": label, "score": prob.item()})
+                for label_index, score in enumerate(predictions_probs[sample_index]):
+                    label = self.config.id2label[label_index]
+                    sample_outputs.append({"label": label, "score": score.item()})
                 outputs.append(sample_outputs)
         else:
             predictions = logits.argmax(1)