|
import streamlit as st |
|
from transformers import AutoModelForTokenClassification |
|
from annotated_text import annotated_text |
|
import numpy as np |
|
import os, joblib |
|
|
|
from utils import get_idxs_from_text |
|
|
|
model = AutoModelForTokenClassification.from_pretrained("CyberPeace-Institute/Cybersecurity-Knowledge-Graph", trust_remote_code=True) |
|
|
|
role_classifiers = {} |
|
folder_path = '/arg_role_models' |
|
for filename in os.listdir(os.getcwd() + folder_path): |
|
if filename.endswith('.joblib'): |
|
file_path = os.getcwd() + os.path.join(folder_path, filename) |
|
clf = joblib.load(file_path) |
|
arg = filename.split(".")[0] |
|
role_classifiers[arg] = clf |
|
|
|
def annotate(name): |
|
tokens = [item["token"] for item in output] |
|
tokens = [token.replace(" ", "") for token in tokens] |
|
text = model.tokenizer.decode([item["id"] for item in output]) |
|
idxs = get_idxs_from_text(text, tokens) |
|
labels = [item[name] for item in output] |
|
|
|
annotated_text_list = [] |
|
last_label = "" |
|
cumulative_tokens = "" |
|
last_id = 0 |
|
for idx, label in zip(idxs, labels): |
|
to_label = label |
|
label_short = to_label.split("-")[1] if "-" in to_label else to_label |
|
if last_label == label_short: |
|
cumulative_tokens += text[last_id : idx["end_idx"]] |
|
last_id = idx["end_idx"] |
|
else: |
|
if last_label != "": |
|
if last_label == "O": |
|
annotated_text_list.append(cumulative_tokens) |
|
else: |
|
annotated_text_list.append((cumulative_tokens, last_label)) |
|
last_label = label_short |
|
cumulative_tokens = idx["word"] |
|
last_id = idx["end_idx"] |
|
if last_label == "O": |
|
annotated_text_list.append(cumulative_tokens) |
|
else: |
|
annotated_text_list.append((cumulative_tokens, last_label)) |
|
annotated_text(annotated_text_list) |
|
|
|
def get_arg_roles(output): |
|
args = [(idx, item["argument"], item["token"]) for idx, item in enumerate(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 = model.tokenizer.decode([item["id"] for item in output[max(0, entity["start"] - 15) : min(len(output), entity["end"] + 15)]]) |
|
entity["context"] = context |
|
|
|
for entity in entities: |
|
if len(model.arg_2_role[entity["label"]]) > 1: |
|
sent_embed = model.embed_model.encode(entity["context"]) |
|
arg_embed = model.embed_model.encode(entity["text"]) |
|
embed = np.concatenate((sent_embed, arg_embed)) |
|
arg_clf = role_classifiers[entity["label"]] |
|
role_id = arg_clf.predict(embed.reshape(1, -1)) |
|
role = model.arg_2_role[entity["label"]][role_id[0]] |
|
entity["role"] = role |
|
else: |
|
entity["role"] = model.arg_2_role[entity["label"]][0] |
|
|
|
for item in output: |
|
item["role"] = "O" |
|
for entity in entities: |
|
for i in range(entity["start"], entity["end"] + 1): |
|
output[i]["role"] = entity["role"] |
|
return output |
|
|
|
st.title("Create Knowledge Graphs from Cyber Incidents") |
|
|
|
text_input = st.text_area("Enter your text here", height=100) |
|
|
|
if text_input: |
|
output = model(text_input) |
|
st.subheader("Event Nuggets") |
|
annotate("nugget") |
|
st.subheader("Event Arguments") |
|
annotate("argument") |
|
st.subheader("Realis of Event Nuggets") |
|
annotate("realis") |
|
output = get_arg_roles(output) |
|
st.subheader("Role of the Event Arguments") |
|
annotate("role") |
|
|