Upload model

modeling.py

@@ -115,7 +115,9 @@ class TotalClassifierModel(PreTrainedModel):
         mask: torch.Tensor | None = None,
         return_logits: bool = False,
         return_as_dict: bool = False,
+        return_as_list: bool = False,
         return_as_df: bool = False,
+        threshold: float = 0.5,  # only used for return_as_list=True
     ) -> torch.Tensor:
         if return_as_df:
             assert (
@@ -138,8 +140,9 @@ class TotalClassifierModel(PreTrainedModel):
             # return raw logits
             return logits
         probas = logits.sigmoid()
+
         if return_as_dict or return_as_df:
-            # list_of_dictionaries
+            # list of dictionaries
             batch_list = []
             for i in range(probas.shape[0]):
                 dict_for_batch = {}
@@ -157,6 +160,21 @@ class TotalClassifierModel(PreTrainedModel):
                 else:
                     batch_list.append(dict_for_batch)
             return batch_list
+
+        if return_as_list:
+            # list of lists
+            batch_list = []
+            for i in range(probas.shape[0]):
+                probas_i = probas[i]
+                batch_list.append(
+                    [
+                        self.index2label[each_class]
+                        for each_class in range(probas_i.shape[1])
+                        if probas_i[:, each_class] >= threshold
+                    ]
+                )
+            return batch_list
+
         return probas
 
     def normalize(self, x: torch.Tensor) -> torch.Tensor: