import streamlit as st
from transformers import pipeline
import pandas as pd
st.set_page_config(layout="wide")

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

st.title("Table Question Answering using TAPAS")

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)

file_name = st.file_uploader("Upload dataset",type=['csv','xlsx'])

if file_name is not None:
  try:
      df=pd.read_csv(file_name)
  except:
      df = pd.read_excel(file_name)
  df = df.astype(str)

  st.markdown("<p style='font-family:sans-serif;font-size: 0.9rem;'> Data - Top 5 records</p>",unsafe_allow_html = True)
  st.table(df.head(5))

  question = st.text_input('Type your question')
  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)