CyberPeace-Institute
/

Cybersecurity-Knowledge-Graph

Token Classification

Model card Files Files and versions

cpi-connect commited on Sep 27, 2023

Commit

dd93bf0

·

1 Parent(s): 12c5ec6

Upload model

Files changed (1) hide show

model.py +1 -1

model.py CHANGED Viewed

@@ -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}
     def forward_model(self, model, dataloader):
         predicted_label = []

             argument_preds[idx] = argument_pred
             realis_preds[idx] = realis_pred
+        return {"nugget" : nugget_pred, "argument" : argument_preds, "realis" : realis_preds}
     def forward_model(self, model, dataloader):
         predicted_label = []