CyberPeace-Institute/Cybersecurity-Knowledge-Graph-Extraction

1import streamlit as st2from transformers import AutoModelForTokenClassification3from annotated_text import annotated_text4import numpy as np5import os, joblib6 7from utils import get_idxs_from_text8 9model = AutoModelForTokenClassification.from_pretrained("CyberPeace-Institute/Cybersecurity-Knowledge-Graph", trust_remote_code=True)10 11role_classifiers = {}12folder_path = '/arg_role_models'13for filename in os.listdir(os.getcwd() + folder_path):14    if filename.endswith('.joblib'):15        file_path = os.getcwd() + os.path.join(folder_path, filename)16        clf = joblib.load(file_path)17        arg = filename.split(".")[0]18        role_classifiers[arg] = clf19 20def annotate(name):21    tokens = [item["token"] for item in output]22    tokens = [token.replace(" ", "") for token in tokens]23    text = model.tokenizer.decode([item["id"] for item in output])24    idxs = get_idxs_from_text(text, tokens)25    labels = [item[name] for item in output]26 27    annotated_text_list = []28    last_label = ""29    cumulative_tokens = "" 30    last_id = 031    for idx, label in zip(idxs, labels):32        to_label = label33        label_short = to_label.split("-")[1] if "-" in to_label else to_label34        if last_label == label_short:35            cumulative_tokens += text[last_id : idx["end_idx"]]36            last_id = idx["end_idx"]37        else:38            if last_label != "":39                if last_label == "O":40                    annotated_text_list.append(cumulative_tokens)41                else:42                    annotated_text_list.append((cumulative_tokens, last_label))43            last_label = label_short44            cumulative_tokens = idx["word"]45            last_id = idx["end_idx"]46    if last_label == "O":47        annotated_text_list.append(cumulative_tokens)48    else:  49        annotated_text_list.append((cumulative_tokens, last_label))50    annotated_text(annotated_text_list)51 52def get_arg_roles(output):53    args = [(idx, item["argument"], item["token"]) for idx, item in enumerate(output) if item["argument"]!= "O"]54        55    entities = []56    current_entity = None57    for position, label, token in args:58        if label.startswith('B-'):59            if current_entity is not None:60                entities.append(current_entity)61            current_entity = {'label': label[2:], 'text': token.replace(" ", ""), 'start': position, 'end': position}62        elif label.startswith('I-'):63            if current_entity is not None:64                current_entity['text'] += ' ' + token.replace(" ", "")65                current_entity['end'] = position66    for entity in entities:67        context = model.tokenizer.decode([item["id"] for item in output[max(0, entity["start"] - 15) : min(len(output), entity["end"] + 15)]])68        entity["context"] = context69    70    for entity in entities:71        if len(model.arg_2_role[entity["label"]]) > 1:72            sent_embed = model.embed_model.encode(entity["context"])73            arg_embed = model.embed_model.encode(entity["text"])74            embed = np.concatenate((sent_embed, arg_embed))75            arg_clf = role_classifiers[entity["label"]]76            role_id = arg_clf.predict(embed.reshape(1, -1))77            role = model.arg_2_role[entity["label"]][role_id[0]]78            entity["role"] = role79        else:80            entity["role"] = model.arg_2_role[entity["label"]][0]81    82    for item in output:83        item["role"] = "O"84    for entity in entities:85        for i in range(entity["start"], entity["end"] + 1):86            output[i]["role"] = entity["role"]87    return output88 89st.title("Create Knowledge Graphs from Cyber Incidents")90 91text_input = st.text_area("Enter your text here", height=100)92 93if text_input or st.button('Apply'):94    output = model(text_input)95    st.subheader("Event Nuggets")96    annotate("nugget")97    st.subheader("Event Arguments")98    annotate("argument")99    st.subheader("Realis of Event Nuggets")100    annotate("realis")101    output = get_arg_roles(output)102    st.subheader("Role of the Event Arguments")103    annotate("role")104