import warnings

warnings.filterwarnings("ignore", category=UserWarning, module="torchvision")
warnings.filterwarnings("ignore", category=FutureWarning, module="transformers")
warnings.filterwarnings("ignore")
from transformers import pipeline
import pandas as pd
import streamlit as st

# Initialize the pipeline
pipe = pipeline("table-question-answering", model="google/tapas-large-finetuned-wtq")

st.set_page_config(page_title="Table Question Answering", layout="wide")
st.title("📊 Table Question Answering")
st.write("""
Welcome to the **Table Question Answering** app! 
You can upload your own table and ask questions about it.
The model will analyze the table and provide accurate answers.
""")

# Step 1: Table input
st.header("📂 Step 1: Upload your table")
# Option to upload a excel file
uploaded_file = st.file_uploader("Upload your Excel file:", type=["xlsx"])
if uploaded_file is not None:
    try:
        table = pd.read_excel(uploaded_file)
        st.success("Table uploaded successfully!")
        st.write("### Preview of Uploaded Table:")
        st.dataframe(table, use_container_width=True)
    except Exception as e:
        st.error(f"Error reading the file: {e}")
        table = None
else:
    st.info("Please upload an Excel file to proceed.")
    table = None

st.divider()
st.header("❓ Step 2: Ask a Question")
query = st.text_input("Ask a question about the table:")


st.divider()
if st.button("🔍 Get Answer"):
    if table is not None and query:
        table = table.astype(str)  # Ensure all values are strings
        result = pipe(table=table, query=query)
        st.success(f"Answer: {result['answer']}")

    else:
        st.warning("Please upload a table and enter a question before clicking 'Get Answer'.")

st.divider()
st.markdown("**❤️**")