Astral7/NER_Roberta

1import streamlit as st2from transformers import AutoTokenizer, TFAutoModelForTokenClassification3from transformers import pipeline4 5 6st.set_page_config(page_title="Test",7                   layout="wide",8                   initial_sidebar_state="expanded")9st.markdown(10    """11    <style>12    body{13    background-color: #27221e;14    }15    [data-testid="stAppViewContainer"]{16        background-color: #27221e;17    }18    [data-testid="stHeader"]{19        background-color: #27221e;20    }21    [data-testid="stToolbar"]{22        color: ##f3d8ba;23    }24    [data-testid="baseButton-secondary"]{25        background-color: #f2ae72;26        color:#27221e;27    }28    [data-testid="stWidgetLabel"]{29        color:#f2ae7230    }31    [data-testid="baseButton-header"]{32        background-color: #f2ae72;33        color:#27221e;34    }35    [data-testid="baseButton-headerNoPadding"]{36        color:#27221e;37        background-color: #f2ae72;38    }39    [data-testid="stStatusWidget"]{40        color:#27221e;41        background-color: #f2ae72;42    }43    </style>44    """,45    unsafe_allow_html=True46)47st.markdown(f'<h1 style="color:#f3d8ba;font-size:24px;text-align:center;font-size:2.75rem;font-weight:700;font-family:"Source Sans Pro", sans-serif">{"Welcome to the Named Entity Recognition App!⚡"}</h1>', unsafe_allow_html=True)48st.markdown(f'<h1 style="color:#f3d8ba;font-size:24px">{"Token Classification"}</h1>', unsafe_allow_html=True)49input_text=st.text_input("Input: ",key="input")50 51submit=st.button("Compute")52 53@st.cache_resource54def classifier(text):55    # Use a pipeline as a high-level helper56    tokenizer = AutoTokenizer.from_pretrained("FacebookAI/roberta-large",add_prefix_space=True)57    model = TFAutoModelForTokenClassification.from_pretrained("Astral7/roberta-large-finetuned-ner")58    pipe = pipeline("token-classification", model=model,tokenizer=tokenizer )59    return pipe(text)60entities=[] 61 62if input_text:63    entities=classifier(input_text)64 65entity_tag = {66 'B-eve':'EVE',67 'I-eve':'EVE',68 'B-org':'ORG',69 'I-org':'ORG',70 'B-gpe':'GPE',71 'I-gpe':'GPE',72 'B-geo':'GEO',73 'I-geo':'GEO',74 'B-nat':'NAT',75 'I-nat':'NAT',76 'B-per':'PER',77 'I-per':'PER',78 'B-art':'ART',79 'I-art':'ART',80 'B-tim':'TIM',81 'I-tim':'TIM',82}83tag_out_styles = {84 'B-eve':"eve_out",85 'I-eve':'eve_out',86 'B-org':'org_out',87 'I-org':'org_out',88 'B-gpe':'gpe_out',89 'I-gpe':'gpe_out',90 'B-geo':'geo_out',91 'I-geo':'geo_out',92 'B-nat':'nat_out',93 'I-nat':'nat_out',94 'B-per':'per_out',95 'I-per':'per_out',96 'B-art':'art_out',97 'I-art':'art_out',98 'B-tim':'tim_out',99 'I-tim':'tim_out',100}101tag_in_styles = {102 'B-eve':'eve_in',103 'I-eve':'eve_in',104 'B-org':'org_in',105 'I-org':'org_in',106 'B-gpe':'gpe_in',107 'I-gpe':'gpe_in',108 'B-geo':'geo_in',109 'I-geo':'geo_in',110 'B-nat':'nat_in',111 'I-nat':'nat_in',112 'B-per':'per_in',113 'I-per':'per_in',114 'B-art':'art_in',115 'I-art':'art_in',116 'B-tim':'tim_in',117 'I-tim':'tim_in',118}119 120custom_style_tag=f"""121<style>122.fl{{123    width:fit-content;124    display:flex;125    align-items:center;126    background-color:#27221e;127}}128.tag_out{{129border-radius:0.25rem;130padding-left: 0.25rem;131padding-right: 0.25rem;132display:flex;133width:max-content;134align-items:center;135margin-right:0.25rem;136}}137.tag_in{{138font-weight:600;139font-size:.75rem;140line-height: 1rem;141border-radius:0.25rem;142margin-left:0.25rem;143height:18px;144padding-left:0.25rem;145padding-right:0.25rem;146margin-top:0.20rem;147margin-bottom:0;148}}149.org_out{{150color:rgb(17 94 89);151background-color: rgb(204 251 241);152}}153.org_in{{154color: rgb(204 251 241);155background-color:rgb(20 184 166);156}}157.per_out{{158color:rgb(91 33 182);159background-color:rgb(237 233 254);160}}161.per_in{{162color: rgb(237 233 254);163background-color:rgb(139 92 246);164}}165.gpe_out{{166color:rgb(134 25 143);167background-color:rgb(250 232 255);168}}169.gpe_in{{170color:rgb(250 232 255);171background-color: rgb(217 70 239);172}}173.eve_in{{174color: rgb(254 226 226);175background-color:rgb(239 68 68) ;176}}177.eve_out{{178color: rgb(239 68 68);179background-color: rgb(254 226 226);180}}181.nat_in{{182color: rgb(255, 239, 213);183background-color: rgb(255, 165, 0);184}}185.nat_out{{186color: rgb(255, 165, 0);187background-color: rgb(255, 239, 213);188}}189.tim_in{{190color: rgb(224 242 254);191background-color: rgb(14 165 233);192}}193.tim_out{{194color: rgb(14 165 233);195background-color:  rgb(224 242 254);196}}197.art_in{{198color: rgb(255, 240, 245);199background-color: rgb(255, 192, 203);200}}201.art_out{{202color: rgb(255, 192, 203);203background-color: rgb(255, 240, 245);204}}205.geo_out{{206color:rgb(157 23 77);207background-color:rgb(252 231 243);208}}209.geo_in{{210color: rgb(252 231 243);211background-color:rgb(236 72 153);212}}213 214</style>215 216<div class="fl">217{"".join([f"<span class='tag_out {tag_out_styles[entity['entity']]}'>{entity['word']}<p class='tag_in {tag_in_styles[entity['entity']]}'>{entity_tag[entity['entity']]}</p></span>" for entity in entities])}218</div>219"""220 221if submit:222    st.markdown(f"<p style='color:#f3d8ba'>Given Input: {input_text}</p>",unsafe_allow_html=True)223    st.markdown(f"<p style='color:#f3d8ba'>Output:</p>",unsafe_allow_html=True)224    st.markdown(custom_style_tag, unsafe_allow_html=True)225