Update model.py

model.py

@@ -136,20 +136,24 @@ class LidirlCNN(PreTrainedModel):
                 probs = torch.softmax(logits, dim=-1)
         return probs
 
-    def predict(self, inputs, lengths, threshold=0.5):
+    def predict(self, inputs, lengths, threshold=0.5, top_k=None):
         probs = self.__call__(inputs, lengths)
+        if top_k is not None and top_k > 0:
+            top_k_preds = torch.topk(probs, top_k, dim=1)
+            pred_labels = []
+            for pred, prob in zip(top_k_preds.indices, top_k_preds.values):
+                pred_labels.append([(self.labels[p.item()], pr.item()) for (p, pr) in zip(pred, prob)])
+            return pred_labels
         if self.multilabel:
             batch_idx, label_idx = torch.where(probs > threshold)
-            output = []
+            output = [[] for _ in range(len(inputs))]
             for batch, label in zip(batch_idx, label_idx):
-                if len(output) <= batch.item():
-                    output.append([])
                 label_string = self.labels
-                output[-1].append(
+                output[batch.item()].append(
                     (self.labels[label.item()], probs[batch, label])
                 )
         return output
-            
+