Upload model

model.py

@@ -88,43 +88,43 @@ class CybersecurityKnowledgeGraphModel(PreTrainedModel):
             structured_output.extend(batch_output)
         
         
-        args = [(idx, item["argument"], item["token"]) for idx, item in enumerate(structured_output) if item["argument"]!= "O"]
+        # args = [(idx, item["argument"], item["token"]) for idx, item in enumerate(structured_output) if item["argument"]!= "O"]
         
-        entities = []
-        current_entity = None
-        for position, label, token in args:
-            if label.startswith('B-'):
-                if current_entity is not None:
-                    entities.append(current_entity)
-                current_entity = {'label': label[2:], 'text': token.replace(" ", ""), 'start': position, 'end': position}
-            elif label.startswith('I-'):
-                if current_entity is not None:
-                    current_entity['text'] += ' ' + token.replace(" ", "")
-                    current_entity['end'] = position
-
-        for entity in entities:
-            context = self.tokenizer.decode([item["id"] for item in structured_output[max(0, entity["start"] - 15) : min(len(structured_output), entity["end"] + 15)]])
-            entity["context"] = context
+        # entities = []
+        # current_entity = None
+        # for position, label, token in args:
+        #     if label.startswith('B-'):
+        #         if current_entity is not None:
+        #             entities.append(current_entity)
+        #         current_entity = {'label': label[2:], 'text': token.replace(" ", ""), 'start': position, 'end': position}
+        #     elif label.startswith('I-'):
+        #         if current_entity is not None:
+        #             current_entity['text'] += ' ' + token.replace(" ", "")
+        #             current_entity['end'] = position
+
+        # for entity in entities:
+        #     context = self.tokenizer.decode([item["id"] for item in structured_output[max(0, entity["start"] - 15) : min(len(structured_output), entity["end"] + 15)]])
+        #     entity["context"] = context
         
-        for entity in entities:
-            if len(self.arg_2_role[entity["label"]]) > 1:
-                sent_embed = self.embed_model.encode(entity["context"])
-                arg_embed = self.embed_model.encode(entity["text"])
-                embed = np.concatenate((sent_embed, arg_embed))
-
-                arg_clf = self.role_classifiers[entity["label"]]
-                role_id = arg_clf.predict(embed.reshape(1, -1))
-                role = self.arg_2_role[entity["label"]][role_id[0]]
-
-                entity["role"] = role
-            else:
-                entity["role"] = self.arg_2_role[entity["label"]][0]
+        # for entity in entities:
+        #     if len(self.arg_2_role[entity["label"]]) > 1:
+        #         sent_embed = self.embed_model.encode(entity["context"])
+        #         arg_embed = self.embed_model.encode(entity["text"])
+        #         embed = np.concatenate((sent_embed, arg_embed))
+
+        #         arg_clf = self.role_classifiers[entity["label"]]
+        #         role_id = arg_clf.predict(embed.reshape(1, -1))
+        #         role = self.arg_2_role[entity["label"]][role_id[0]]
+
+        #         entity["role"] = role
+        #     else:
+        #         entity["role"] = self.arg_2_role[entity["label"]][0]
         
-        for item in structured_output:
-            item["role"] = "O"
-        for entity in entities:
-            for i in range(entity["start"], entity["end"] + 1):
-                structured_output[i]["role"] = entity["role"]
+        # for item in structured_output:
+        #     item["role"] = "O"
+        # for entity in entities:
+        #     for i in range(entity["start"], entity["end"] + 1):
+        #         structured_output[i]["role"] = entity["role"]
         return structured_output
 
     def forward_model(self, model, dataloader):