| """ |
| The SpEL annotation visualization script. You can use this script as a playground to explore the capabilities and |
| limitations of the SpEL framework. |
| """ |
| import torch |
| from model import SpELAnnotator |
| from data_loader import dl_sa |
| from utils import chunk_annotate_and_merge_to_phrase |
| from candidate_manager import CandidateManager |
| import streamlit as st |
| from annotated_text import annotated_text |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
| @st.cache_resource |
| def load_model(): |
| load_aida_finetuned = True |
| load_full_vocabulary=True |
| candidate_setting = "n" |
| model = SpELAnnotator() |
| model.init_model_from_scratch(device=device) |
| candidates_manager_to_use = CandidateManager(dl_sa.mentions_vocab, |
| is_kb_yago=candidate_setting == "k", |
| is_ppr_for_ned=candidate_setting.startswith("p"), |
| is_context_agnostic=candidate_setting == "pg", |
| is_indexed_for_spans=True) if candidate_setting != "n" else None |
| if load_aida_finetuned and not load_full_vocabulary: |
| model.shrink_classification_head_to_aida(device=device) |
| model.load_checkpoint(None, device=device, load_from_torch_hub=True, finetuned_after_step=3) |
| elif load_aida_finetuned: |
| model.load_checkpoint(None, device=device, load_from_torch_hub=True, finetuned_after_step=4) |
| else: |
| model.load_checkpoint(None, device=device, load_from_torch_hub=True, finetuned_after_step=2) |
| return model, candidates_manager_to_use |
|
|
| annotator, candidates_manager = load_model() |
| st.title("SpEL Prediction Visualization") |
| st.caption('Running the \"[SpEL-base-step3-500K.pt](https://vault.sfu.ca/index.php/s/8nw5fFXdz2yBP5z/download)\" model without any hand-crafted candidate sets. For more information please checkout [SpEL\'s github repository](https://github.com/shavarani/SpEL).') |
| mention = st.text_input("Enter the text:") |
| process_button = st.button("Annotate") |
|
|
| if process_button and mention: |
| phrase_annotations = chunk_annotate_and_merge_to_phrase( |
| annotator, mention, k_for_top_k_to_keep=5, normalize_for_chinese_characters=True) |
| last_step_annotations = [[p.words[0].token_offsets[0][1][0], |
| p.words[-1].token_offsets[-1][1][-1], |
| (dl_sa.mentions_itos[p.resolved_annotation], p.subword_annotations)] |
| for p in phrase_annotations if p.resolved_annotation != 0] |
| if candidates_manager: |
| for p in phrase_annotations: |
| candidates_manager.modify_phrase_annotation_using_candidates(p, mention) |
| if last_step_annotations: |
| anns = sorted([(l_ann[0], l_ann[1], l_ann[2][0]) for l_ann in last_step_annotations], key=lambda x: x[0]) |
| begin = 0 |
| last_char = len(mention) |
| anns_pointer = 0 |
| processed_anns = [] |
| anno_text = [] |
| while begin < last_char: |
| if anns_pointer == len(anns): |
| processed_anns.append((begin, last_char, "O")) |
| anno_text.append(mention[begin: last_char]) |
| begin = last_char |
| continue |
| first_unprocessed_annotation = anns[anns_pointer] |
| if first_unprocessed_annotation[0] > begin: |
| processed_anns.append((begin, first_unprocessed_annotation[0], "O")) |
| anno_text.append(mention[begin: first_unprocessed_annotation[0]]) |
| begin = first_unprocessed_annotation[0] |
| else: |
| processed_anns.append(first_unprocessed_annotation) |
| anns_pointer += 1 |
| begin = first_unprocessed_annotation[1] |
| anno_text.append((mention[first_unprocessed_annotation[0]: first_unprocessed_annotation[1]], first_unprocessed_annotation[2])) |
| annotated_text(anno_text) |
| else: |
| annotated_text(mention) |
|
|
