Spaces:
Paused
Paused
| import streamlit as st | |
| from spacy import displacy | |
| from Model.NER.VLSP2021.Predict_Ner import ViTagger | |
| import re | |
| import torch | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| def process_text(text): | |
| # Loại bỏ dấu cách thừa và dấu cách ở đầu và cuối văn bản | |
| processed_text = re.sub(r'\s+', ' ', text.strip()) | |
| return processed_text | |
| def show_ner(): | |
| st.sidebar.title('Datasets') | |
| dataset = st.sidebar.selectbox("Datasets", ("VLSP2016", "VLSP2021")) | |
| st.header("NER") | |
| text = st.text_area("Enter your text for NER:", height=300) | |
| text = process_text(text) | |
| if st.button("Process NER"): | |
| if dataset == "VLSP2021": | |
| tagger = ViTagger(model_path='E:/demo_datn/pythonProject1/Model/NER/VLSP2021/best_model.pt') | |
| a = text | |
| b = tagger.extract_entity_doc(a) | |
| # Danh sách các từ và nhãn NER | |
| # words_and_labels = [('Dự', 'O'), ('báo', 'O'), ('thời', 'O'), ('tiết', 'O'), ('2 ngày', 'DATETIME-DATERANGE'), ('tới', 'O'), ('Nam Bộ', 'LOCATION-GPE'), ('có', 'O'), ('mưa', 'O'), ('dông', 'O')] | |
| words_and_labels = b | |
| # Tạo danh sách từ | |
| words = [word for word, _ in words_and_labels] | |
| # Tạo danh sách thực thể và nhãn cho mỗi từ, loại bỏ nhãn 'O' | |
| entities = [{'start': sum(len(word) + 1 for word, _ in words_and_labels[:i]), | |
| 'end': sum(len(word) + 1 for word, _ in words_and_labels[:i + 1]), 'label': label} for | |
| i, (word, label) | |
| in enumerate(words_and_labels) if label != 'O'] | |
| # print(entities) | |
| # Render the visualization without color for 'O' labels | |
| html = displacy.render( | |
| {"text": " ".join(words), "ents": entities, "title": None}, | |
| style="ent", | |
| manual=True, | |
| options={"colors": {"DATETIME-DATERANGE": "#66c2ff", | |
| "LOCATION-GPE": "#ffcc99", | |
| "O": None, # Màu cho nhãn 'O' | |
| "QUANTITY-NUM": "#ffdf80", | |
| "EVENT-CUL": "#bfbfbf", | |
| "DATETIME": "#80ff80", | |
| "PERSONTYPE": "#ff80ff", | |
| "PERSON": "#bf80ff", | |
| "QUANTITY-PER": "#80cccc", | |
| "ORGANIZATION": "#ff6666", | |
| "LOCATION-GEO": "#66cc66", | |
| "LOCATION-STRUC": "#cccc66", | |
| "PRODUCT-COM": "#ffff66", | |
| "DATETIME-DATE": "#66cccc", | |
| "QUANTITY-DIM": "#6666ff", | |
| "PRODUCT": "#cc6666", | |
| "QUANTITY": "#6666cc", | |
| "DATETIME-DURATION": "#9966ff", | |
| "QUANTITY-CUR": "#ff9966", | |
| "DATETIME-TIME": "#cdbf93", | |
| "QUANTITY-TEM": "#cc9966", | |
| "DATETIME-TIMERANGE": "#cc8566", | |
| "EVENT-GAMESHOW": "#8c8c5a", | |
| "QUANTITY-AGE": "#70db70", | |
| "QUANTITY-ORD": "#e699ff", | |
| "PRODUCT-LEGAL": "#806699", | |
| "LOCATION": "#993366", | |
| "ORGANIZATION-MED": "#339933", | |
| "URL": "#ff4d4d", | |
| "PHONENUMBER": "#99cc99", | |
| "ORGANIZATION-SPORTS": "#6666ff", | |
| "EVENT-SPORT": "#ffff80", | |
| "SKILL": "#b38f66", | |
| "EVENT-NATURAL": "#ff9966", | |
| "ADDRESS": "#cc9966", | |
| "IP": "#b38f66", | |
| "EMAIL": "#cc8566", | |
| "ORGANIZATION-STOCK": "#666633", | |
| "DATETIME-SET": "#70db70", | |
| "PRODUCT-AWARD": "#e699ff", | |
| "MISCELLANEOUS": "#806699", | |
| "LOCATION-GPE-GEO": "#99ffff"}} | |
| ) | |
| # print(html) | |
| st.markdown(html, unsafe_allow_html=True) | |
| elif dataset == "VLSP2016": | |
| tagger = ViTagger(model_path='E:/demo_datn/pythonProject1/Model/NER/VLSP2016/best_model.pt') | |
| a = text | |
| b = tagger.extract_entity_doc(a) | |
| # Danh sách các từ và nhãn NER | |
| # words_and_labels = [('Dự', 'O'), ('báo', 'O'), ('thời', 'O'), ('tiết', 'O'), ('2 ngày', 'DATETIME-DATERANGE'), ('tới', 'O'), ('Nam Bộ', 'LOCATION-GPE'), ('có', 'O'), ('mưa', 'O'), ('dông', 'O')] | |
| words_and_labels = b | |
| # Tạo danh sách từ | |
| words = [word for word, _ in words_and_labels] | |
| # Tạo danh sách thực thể và nhãn cho mỗi từ, loại bỏ nhãn 'O' | |
| entities = [{'start': sum(len(word) + 1 for word, _ in words_and_labels[:i]), | |
| 'end': sum(len(word) + 1 for word, _ in words_and_labels[:i + 1]), 'label': label} for | |
| i, (word, label) | |
| in enumerate(words_and_labels) if label != 'O'] | |
| # print(entities) | |
| # Render the visualization without color for 'O' labels | |
| html = displacy.render( | |
| {"text": " ".join(words), "ents": entities, "title": None}, | |
| style="ent", | |
| manual=True, | |
| options={"colors": {"MISC": "#806699", | |
| "ORG": "#ff6666", | |
| "LOC": "#66cc66", | |
| "PER": "#bf80ff", | |
| "O": None}} | |
| ) | |
| # print(html) | |
| st.markdown(html, unsafe_allow_html=True) | |
| # Sử dụng widget st.html để hiển thị HTML | |
| # Hiển thị văn bản đã nhập | |
| # st.write("Văn bản đã nhập:", text) | |