Upload model

event_arg_predict.py

@@ -3,11 +3,11 @@ from annotated_text import annotated_text
 import torch
 from torch.utils.data import DataLoader
 
-from cybersecurity_knowledge_graph.args_model_utils import tokenize_and_align_labels_with_pos_ner_dep, find_nearest_nugget_features, find_dep_depth
-from cybersecurity_knowledge_graph.nugget_model_utils import CustomRobertaWithPOS
-from cybersecurity_knowledge_graph.utils import get_content, get_event_nugget, get_idxs_from_text, get_entity_from_idx, list_of_pos_tags, event_args_list
+from .args_model_utils import tokenize_and_align_labels_with_pos_ner_dep, find_nearest_nugget_features, find_dep_depth
+from .nugget_model_utils import CustomRobertaWithPOS
+from .utils import get_content, get_event_nugget, get_idxs_from_text, get_entity_from_idx, list_of_pos_tags, event_args_list
 
-from cybersecurity_knowledge_graph.event_nugget_predict import get_event_nuggets
+from .event_nugget_predict import get_event_nuggets
 import spacy
 from transformers import AutoTokenizer
 from datasets import load_dataset, Features, ClassLabel, Value, Sequence, Dataset
@@ -37,7 +37,7 @@ device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cp
 model_checkpoint = "ehsanaghaei/SecureBERT"
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=True)
 
-from cybersecurity_knowledge_graph.args_model_utils import CustomRobertaWithPOS as ArgumentModel
+from .args_model_utils import CustomRobertaWithPOS as ArgumentModel
 model_nugget = ArgumentModel(num_classes=43)
 model_nugget.load_state_dict(torch.load("cybersecurity_knowledge_graph/argument_model_state_dict.pth", map_location=device)) 
 model_nugget.eval()

event_nugget_predict.py

@@ -3,9 +3,9 @@ from annotated_text import annotated_text
 import torch
 from torch import nn
 from torch.utils.data import DataLoader
-from cybersecurity_knowledge_graph.nugget_model_utils import CustomRobertaWithPOS as NuggetModel
-from cybersecurity_knowledge_graph.nugget_model_utils import tokenize_and_align_labels_with_pos_ner_dep, find_nearest_nugget_features, find_dep_depth
-from cybersecurity_knowledge_graph.utils import get_idxs_from_text, event_nugget_list
+from .nugget_model_utils import CustomRobertaWithPOS as NuggetModel
+from .nugget_model_utils import tokenize_and_align_labels_with_pos_ner_dep, find_nearest_nugget_features, find_dep_depth
+from .utils import get_idxs_from_text, event_nugget_list
 import spacy
 from transformers import AutoTokenizer
 from datasets import load_dataset, Features, ClassLabel, Value, Sequence, Dataset

event_realis_predict.py

@@ -3,12 +3,12 @@ import spacy
 import torch
 from torch.utils.data import DataLoader
 from transformers import AutoTokenizer
-from cybersecurity_knowledge_graph.utils import get_idxs_from_text
+from .utils import get_idxs_from_text
 import streamlit as st
 from annotated_text import annotated_text
-from cybersecurity_knowledge_graph.nugget_model_utils import CustomRobertaWithPOS
-from cybersecurity_knowledge_graph.event_nugget_predict import get_event_nuggets
-from cybersecurity_knowledge_graph.realis_model_utils import get_entity_for_realis_from_idx, tokenize_and_align_labels_with_pos_ner_realis
+from .nugget_model_utils import CustomRobertaWithPOS
+from .event_nugget_predict import get_event_nuggets
+from .realis_model_utils import get_entity_for_realis_from_idx, tokenize_and_align_labels_with_pos_ner_realis
 from datasets import load_dataset, Features, ClassLabel, Value, Sequence, Dataset
 
 event_nugget_list = ['B-Phishing',
@@ -49,7 +49,7 @@ device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cp
 model_checkpoint = "ehsanaghaei/SecureBERT"
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=True)
 
-from cybersecurity_knowledge_graph.realis_model_utils import CustomRobertaWithPOS as RealisModel
+from .realis_model_utils import CustomRobertaWithPOS as RealisModel
 model_realis = RealisModel(num_classes_realis=4)
 model_realis.load_state_dict(torch.load("cybersecurity_knowledge_graph/realis_model_state_dict.pth", map_location=device))  
 model_realis.eval()