Update model.py

model.py

@@ -110,6 +110,10 @@ class LidirlCNN(PreTrainedModel):
         self.multilabel = config.multilabel
         self.monte_carlo = config.montecarlo_layer
 
+        self.labels = ["" for _ in config.labels]
+        for key, value in config.labels.items():
+            self.labels[value] = key
+
 
     def forward(self, inputs, lengths):
         inputs = inputs[:, :self.max_length]
@@ -123,9 +127,29 @@ class LidirlCNN(PreTrainedModel):
         return projection
 
     def __call__(self, inputs, lengths):
-        logits = self.forward(inputs, lengths)
+        # this is inference only model
+        with torch.no_grad():
+            logits = self.forward(inputs, lengths)
+            if self.multilabel:
+                probs = torch.sigmoid(logits)
+            else:
+                probs = torch.softmax(logits, dim=-1)
+        return probs
+
+    def predict(self, inputs, lengths, threshold=0.5):
+        probs = self.__call__(inputs, lengths)
         if self.multilabel:
-            probs = torch.sigmoid(logits)
-        else:
-            probs = torch.softmax(logits, dim=-1)
-        return probs
+            batch_idx, label_idx = torch.where(probs > threshold)
+            output = []
+            for batch, label in zip(batch_idx, label_idx):
+                if len(output) < batch.item():
+                    output.append([])
+                label_string = self.labels
+                output[-1].append(
+                    (self.labels[label.item()], probs[batch, label])
+                )
+        return output
+            
+                
+                    
+