Upload model

model.py

@@ -52,7 +52,7 @@ class CybersecurityKnowledgeGraphModel(PreTrainedModel):
             argument_preds[idx] = argument_pred
             realis_preds[idx] = realis_pred
 
-        return {"nugget" : nugget_pred, "argument" : argument_pred, "realis" : realis_pred}
+        return {"nugget" : nugget_pred, "argument" : argument_preds, "realis" : realis_preds}
 
     def forward_model(self, model, dataloader):
         predicted_label = []