Linhz/ViMNer

1import streamlit as st2from spacy import displacy3from Model.NER.VLSP2021.Predict_Ner import ViTagger4import re5from thunghiemxuly import save_uploaded_image,convert_text_to_txt,add_string_to_txt6 7import os8from transformers import AutoTokenizer, BertConfig9from Model.MultimodelNER.VLSP2016.train_umt_2016 import load_model,predict10from Model.MultimodelNER.Ner_processing import format_predictions,process_predictions,combine_entities,remove_B_prefix,combine_i_tags11 12from Model.MultimodelNER.predict import get_test_examples_predict13from Model.MultimodelNER import resnet as resnet14from Model.MultimodelNER.resnet_utils import myResnet15import torch16import numpy as np17from Model.MultimodelNER.VLSP2016.dataset_roberta import MNERProcessor_201618 19 20CONFIG_NAME = 'bert_config.json'21WEIGHTS_NAME = 'pytorch_model.bin'22device = torch.device("cuda" if torch.cuda.is_available() else "cpu")23 24 25net = getattr(resnet, 'resnet152')()26net.load_state_dict(torch.load(os.path.join('Model/Resnet/', 'resnet152.pth')))27encoder = myResnet(net, True, device)28def process_text(text):29    # Loại bỏ dấu cách thừa và dấu cách ở đầu và cuối văn bản30    processed_text = re.sub(r'\s+', ' ', text.strip())31    return processed_text32 33 34 35def show_mner_2016():36    multimodal_text = st.text_area("Enter your text for MNER:", height=300)37    multimodal_text = process_text(multimodal_text)  # Xử lý văn bản38    image = st.file_uploader("Upload an image (only jpg):", type=["jpg"])39    if st.button("Process Multimodal NER"):40            save_image = 'Model/MultimodelNER/VLSP2016/Image'41            save_txt = 'Model/MultimodelNER/VLSP2016/Filetxt/test.txt'42            image_name = image.name43            save_uploaded_image(image, save_image)44            convert_text_to_txt(multimodal_text, save_txt)45            add_string_to_txt(image_name, save_txt)46            st.image(image, caption="Uploaded Image", use_column_width=True)47 48            bert_model='vinai/phobert-base-v2'49            output_dir='Model/MultimodelNER/VLSP2016/best_model'50            output_model_file = os.path.join(output_dir, WEIGHTS_NAME)51            output_encoder_file = os.path.join(output_dir, "pytorch_encoder.bin")52            processor = MNERProcessor_2016()53            label_list = processor.get_labels()54            auxlabel_list = processor.get_auxlabels()55            num_labels = len(label_list) + 156            auxnum_labels = len(auxlabel_list) + 157            trans_matrix = np.zeros((auxnum_labels, num_labels), dtype=float)58            trans_matrix[0, 0] = 1  # pad to pad59            trans_matrix[1, 1] = 1  # O to O60            trans_matrix[2, 2] = 0.25  # B to B-MISC61            trans_matrix[2, 4] = 0.25  # B to B-PER62            trans_matrix[2, 6] = 0.25  # B to B-ORG63            trans_matrix[2, 8] = 0.25  # B to B-LOC64            trans_matrix[3, 3] = 0.25  # I to I-MISC65            trans_matrix[3, 5] = 0.25  # I to I-PER66            trans_matrix[3, 7] = 0.25  # I to I-ORG67            trans_matrix[3, 9] = 0.25  # I to I-LOC68            trans_matrix[4, 10] = 1  # X to X69            trans_matrix[5, 11] = 1  # [CLS] to [CLS]70            trans_matrix[6, 12] = 171            tokenizer = AutoTokenizer.from_pretrained(bert_model, do_lower_case=False)72            model_umt, encoder_umt = load_model(output_model_file, output_encoder_file, encoder,num_labels,auxnum_labels)73            eval_examples = get_test_examples_predict('Model/MultimodelNER/VLSP2016/Filetxt/')74 75            y_pred, a = predict(model_umt, encoder_umt, eval_examples, tokenizer, device,save_image,trans_matrix)76            formatted_output = format_predictions(a, y_pred[0])77            final = process_predictions(formatted_output)78            final2 = combine_entities(final)79            final3 = remove_B_prefix(final2)80            final4 = combine_i_tags(final3)81            words_and_labels = final482            # Tạo danh sách từ83            words = [word for word, _ in words_and_labels]84            # Tạo danh sách thực thể và nhãn cho mỗi từ, loại bỏ nhãn 'O'85            entities = [{'start': sum(len(word) + 1 for word, _ in words_and_labels[:i]),86                         'end': sum(len(word) + 1 for word, _ in words_and_labels[:i + 1]), 'label': label} for87                        i, (word, label)88                        in enumerate(words_and_labels) if label != 'O']89            # print(entities)90 91            # Render the visualization without color for 'O' labels92            html = displacy.render(93                {"text": " ".join(words), "ents": entities, "title": None},94                style="ent",95                manual=True,96                options={"colors": {"MISC": "#806699",97                                    "ORG": "#ff6666",98                                    "LOC": "#66cc66",99                                    "PER": "#bf80ff",100                                    "O": None}}101            )102            # print(html)103            st.markdown(html, unsafe_allow_html=True)104 105 106###Ví dụ 1 : Một trận hỗn chiến đã xảy ra tại trận đấu khúc côn cầu giữa  Penguins và Islanders ở Mỹ (image:penguin)