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 plotly.express as px

st.subheader("Table Question Answering (QA)", divider="blue")

# 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 Table Question Answering (QA) App**")
    expander.write('''
    
    **Supported File Formats**
    This app accepts files in .csv and .xlsx formats.
    
    **How to Use**
    Upload your file first. Then, type your question into the text area provided and click the 'Retrieve your answer' button.
    
    **Usage Limits**
    You can ask up to 5 questions. 
    
    **Subscription Management**
    This demo app offers a one-day subscription, expiring after 24 hours. If you are interested in building your own Table Question Answering (QA) 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 in five business days. If you wish to delete your Account with us, please contact us at info@nlpblogs.com
    
    **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 has errors or date values.
    For any errors or inquiries, please contact us at info@nlpblogs.com
    
''')
    
    
# count attempts based on questions
if 'question_attempts' not in st.session_state:
    st.session_state['question_attempts'] = 0

max_attempts = 5

# upload file
upload_file = st.file_uploader("Upload your file. Accepted file formats include: .csv, .xlsx", type=['csv', 'xlsx'])


if upload_file is not None:
    file_extension = upload_file.name.split('.')[-1].lower()
    if file_extension == 'csv':
        try:
            df = pd.read_csv(upload_file, na_filter=False)
            if df.isnull().values.any():
                st.error("Error: The CSV file contains missing values.")
                st.stop()
            else:
                new_columns = [f'column_{i+1}' for i in range(len(df.columns))]
                df.columns = new_columns
                st.dataframe(df, key="csv_dataframe")
                all_columns = df.columns.tolist()
                st.subheader("Select columns for the Tree Map", divider="blue")
                parent_column = st.selectbox("Select the parent column:", all_columns)
                value_column = st.selectbox("Select the value column:", all_columns)
                
            if parent_column and value_column:
                if parent_column == value_column:
                    st.warning("Warning: You have selected the same column for both the parent and value. This might not produce a meaningful treemap.")
                elif parent_column and value_column:
                    path_columns = [px.Constant("all"), parent_column, value_column]
                    fig = px.treemap(df,
                         path=path_columns)

                    fig.update_layout(margin=dict(t=50, l=25, r=25, b=25))
                    st.subheader("Tree map", divider="red")
                    st.plotly_chart(fig)
                    
                    st.write("_number of rows_", df.shape[0])
                    st.write("_number of columns_", df.shape[1])
                    st.session_state.df = df  
        except pd.errors.ParserError:
            st.error("Error: The CSV file is not readable or is incorrectly formatted.")
            st.stop()
        except UnicodeDecodeError:
            st.error("Error: The CSV file could not be decoded.")
            st.stop()
        except Exception as e:
            st.error(f"An unexpected error occurred while reading CSV: {e}")
            st.stop()
    elif file_extension == 'xlsx':
        try:
            df = pd.read_excel(upload_file, na_filter=False)
            
            if df.isnull().values.any():
                st.error("Error: The Excel file contains missing values.")
                st.stop()
            else:
                new_columns = [f'column_{i+1}' for i in range(len(df.columns))]
                df.columns = new_columns
                st.write(df.columns)
                st.dataframe(df, key="excel_dataframe")
                
                
                st.write("_number of rows_", df.shape[0])
                st.write("_number of columns_", df.shape[1])
                st.session_state.df = df  
        except ValueError:
            st.error("Error: The Excel file is not readable or is incorrectly formatted.")
            st.stop()
        except Exception as e:
            st.error(f"An unexpected error occurred while reading Excel: {e}")
            st.stop()
    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 'df' in st.session_state: 
        df = st.session_state.df
    
        st.session_state.fernet_token = generate_fernet_token(key, df.to_json())
    else:
        st.stop()

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

if error_streamlit:
    st.warning("Please press Request Authorization. Please note that a file should be uploaded before you press Request Authorization.")
    if st.button("Request Authorization"):
        if 'df' in st.session_state: 
            df = st.session_state.df
            st.session_state.fernet_token = generate_fernet_token(key, df.to_json())
            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()
        if error_streamlit:
            st.error("Please upload a file.")
            st.stop()
        else:
            try:
                df = pd.read_json(decrypted_data_streamlit)
            except Exception as e:
                st.error(f"Error decoding data: {e}")
                st.stop()
    else:
        st.error(f"Your authorization has expired: {error_streamlit}")
        st.stop()

st.divider()

# ask question
def clear_question():
    st.session_state["question"] = ""

question = st.text_input("Type your question here and then press **Retrieve your answer**:", key="question")
st.button("Clear question", on_click=clear_question)

#retrive answer
if st.button("Retrieve your answer"):
    if st.session_state['question_attempts'] >= max_attempts:
        st.error(f"You have asked {max_attempts} questions. Maximum question attempts reached.")
        st.stop()
    st.session_state['question_attempts'] += 1
    if error_streamlit:
        st.warning("Please enter a question before retrieving the answer.")
    else:
        with st.spinner('Wait for it...'):
            time.sleep(2)
            if df is not None:
                tqa = pipeline(task="table-question-answering", model="microsoft/tapex-large-finetuned-wtq")
                st.write(tqa(table=df, query=question)['answer'])

st.divider()
st.write(f"Number of questions asked: {st.session_state['question_attempts']}/{max_attempts}")