app.py

import streamlit as st
from cryptography.fernet import Fernet
import time
import pandas as pd
import io
from transformers import pipeline
from streamlit_extras.stylable_container import stylable_container
import json

import nltk

import plotly.express as px
from PyPDF2 import PdfReader
import docx
import zipfile

from gliner import GLiNER



st.subheader("Named Entity Recognition (NER)", divider="red")

# generate Fernet key
if 'fernet_key' not in st.session_state:
    st.session_state.fernet_key = Fernet.generate_key()

key = st.session_state.fernet_key


# function for generating and validating fernet key
def generate_fernet_token(key, data):
    fernet = Fernet(key)
    token = fernet.encrypt(data.encode())
    return token

def validate_fernet_token(key, token, ttl_seconds):
    
    fernet = Fernet(key)
    try:
        decrypted_data = fernet.decrypt(token, ttl=ttl_seconds).decode()
        return decrypted_data, None
    except Exception as e:
        return None, f"Expired token: {e}"


# sidebar   
with st.sidebar:
    with stylable_container(
        key="test_button",
        css_styles="""
        button {
            background-color: yellow;
            border: 1px solid black;
            padding: 5px;
            color: black;
        }
        """,
    ):
        st.button("DEMO APP")
   

    expander = st.expander("**Important notes on the Demo Named Entity Recognition (NER) App**")
    expander.write('''
    
    **Supported File Formats**
    This app accepts files in .pdf and .docx formats.
    
    **How to Use**
    Upload your file first. Then, click the 'Results' button.
    
    **Usage Limits**
    You can request results up to 5 times. 
    
     **Subscription Management**
    This demo app offers a one-day subscription, expiring after 24 hours. If you are interested in building your own Named Entity Recognition (NER) Web App, we invite you to explore our NLP Web App Store on our website. You can select your desired features, place your order, and we will deliver your custom app within five business days. If you wish to delete your Account with us, please contact us at [email protected]
    
    **Authorization**
    For security purposes, your authorization access expires hourly. To restore access, click the "Request Authorization" button. 
    
    **Customization**
    To change the app's background color to white or black, click the three-dot menu on the right-hand side of your app, go to Settings and then Choose app theme, colors and fonts.
    
    **File Handling and Errors**
    The app may display an error message if your file is corrupt, or has other errors.
    
    
    For any errors or inquiries, please contact us at [email protected]
    
''')
    
    
       
# count attempts based on file upload
if 'file_upload_attempts' not in st.session_state:
    st.session_state['file_upload_attempts'] = 0

max_attempts = 5

# upload file
upload_file = st.file_uploader("Upload your file. Accepted file formats include: .pdf, .docx", type=['pdf', 'docx'])
text = None
df = None

if upload_file is not None:
    
    file_extension = upload_file.name.split('.')[-1].lower()
    if file_extension == 'pdf':
        try:
            pdf_reader = PdfReader(upload_file)
            text = ""
            for page in pdf_reader.pages:
                text += page.extract_text()
            st.write(text)
        except Exception as e:
            st.error(f"An error occurred while reading PDF: {e}")
    elif file_extension == 'docx':
        try:
            doc = docx.Document(upload_file)
            text = "\n".join([para.text for para in doc.paragraphs])
            st.write(text)
        except Exception as e:
            st.error(f"An error occurred while reading docx: {e}")
    else:
        st.warning("Unsupported file type.")
    
        st.stop()


        

# generate and validate Fernet token for the current file
if 'fernet_token' not in st.session_state:
    if text is not None:
        st.session_state.fernet_token = generate_fernet_token(key, text)
    else:
        st.stop()

decrypted_data_streamlit, error_streamlit = validate_fernet_token(key, st.session_state.fernet_token, ttl_seconds=3600)

if error_streamlit:
    if text is not None:
        st.warning("Please press Request Authorization.")
        if st.button("Request Authorization"):
            st.session_state.fernet_token = generate_fernet_token(key, text)
            st.success("Authorization granted")
            decrypted_data_streamlit, error_streamlit = validate_fernet_token(key, st.session_state.fernet_token, ttl_seconds=3600)
        if error_streamlit:
            st.error(f"Your authorization has expired: {error_streamlit}")
            st.stop()
        

st.divider()



#retrieve answer
if st.button("Results"):
    if st.session_state['file_upload_attempts'] >= max_attempts:
        st.error(f"You have requested results {max_attempts} times. You have reached your daily request limit.")
        st.stop()
    st.session_state['file_upload_attempts'] += 1
    if error_streamlit:
        st.warning("Please upload a file before retrieving the results.")
    else:
        with st.spinner('Wait for it...'):
            time.sleep(2)
            model = GLiNER.from_pretrained("xomad/gliner-model-merge-large-v1.0")
            labels = ["person", "location", "country", "city", "organization", "time", "date", "product", "event name", "money", "affiliation", "ordinal value", "percent value", "position"]
            entities = model.predict_entities(text, labels)
            df = pd.DataFrame(entities)
            
            properties = {"border": "2px solid gray", "color": "blue", "font-size": "16px"}
            df_styled = df.style.set_properties(**properties)
            st.dataframe(df_styled)
            if df is not None:
                fig = px.treemap(df, path=[px.Constant("all"), 'text', 'label'],
                values='score', color='label')
                fig.update_layout(margin = dict(t=50, l=25, r=25, b=25))
                st.plotly_chart(fig)


            
                
            

            dfa = pd.DataFrame(
                data={
                    
                    'text': ['entity extracted from file'], 'score': ['accuracy score'], 'label': ['label assigned to the extracted entity'],
                    'start': ['index of the start of the corresponding entity'],
                    'end': ['index of the end of the corresponding entity'],
                })
        
            buf = io.BytesIO()
            with zipfile.ZipFile(buf, "w") as myzip:
                myzip.writestr("Summary of the results.csv", df.to_csv(index=False))
                myzip.writestr("Glossary of labels.csv", dfa.to_csv(index=False))
                    
            with stylable_container(
                key="download_button",
                css_styles="""button { background-color: yellow; border: 1px solid black; padding: 5px; color: black; }""",
            ):
                st.download_button(
                        label="Download zip file",
                        data=buf.getvalue(),
                        file_name="zip file.zip",
                        mime="application/zip",
                )
                    





            


st.divider()
st.write(f"Number of times you requested results: {st.session_state['file_upload_attempts']}/{max_attempts}")