import streamlit as st
from st_aggrid import AgGrid
import pandas as pd 
# from PIL import Image
from transformers import pipeline
st.set_page_config(layout="wide")

# im = Image.open("ai-favicon.png")
# st.set_page_config(page_title="Table Summarization",
#     page_icon=im,layout='wide') 


style = '''
    <style>
        header {visibility: hidden;}
        div.block-container {padding-top:4rem;}
        section[data-testid="stSidebar"] div:first-child {
        padding-top: 0;
    }
     .font {                                          
    text-align:center;
    font-family:sans-serif;font-size: 1.25rem;}
    </style>
'''
st.markdown(style, unsafe_allow_html=True)

st.markdown('<p style="font-family:sans-serif;font-size: 1.9rem;">Table Question Answering using TAPAS</p>', unsafe_allow_html=True)
st.markdown("<p style='font-family:sans-serif;font-size: 0.9rem;'>Pre-trained TAPAS model runs on max 64 rows and 32 columns data. Make sure the file data doesn't exceed these dimensions.</p>", unsafe_allow_html=True)

tqa = pipeline(task="table-question-answering", 
                    model="google/tapas-base-finetuned-wtq")


# st.sidebar.image("ai-logo.png",width=200)
# with open('data.csv', 'rb') as f:
#         st.sidebar.download_button('Download sample data', f, file_name='Sample Data.csv')
file_name = st.sidebar.file_uploader("Upload file:", type=['csv','xlsx'])

if file_name is None:
    st.markdown('<p class="font">Please upload an excel or csv file </p>', unsafe_allow_html=True)
    # st.image("loader.png")

else:
    try:
        df=pd.read_csv(file_name)
    except:
        df = pd.read_excel(file_name)
         
    grid_response = AgGrid(
        df.head(5),
        columns_auto_size_mode='FIT_CONTENTS',
        editable=True, 
        height=300, 
        width='100%',
        )

    question = st.text_input('Type your question')
    df = df.astype(str)
    
    with st.spinner():
        if(st.button('Answer')):
            answer = tqa(table=df, query=question,truncation=True)
            st.markdown("<p style='font-family:sans-serif;font-size: 0.9rem;'> Results </p>",unsafe_allow_html = True)
            st.success(answer)