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")

# Define the table data
data = {
    "year": [1896, 1900, 1904, 2004, 2008, 2012],
    "city": ["athens", "paris", "st. louis", "athens", "beijing", "london"]
}
table = pd.DataFrame(data)
table = table.astype(str)  # Ensure all values are strings

# Streamlit app layout
st.title("Table Question Answering")
st.write("### Input Table")
st.dataframe(table)

# User query
query = st.text_input("Ask a question about the table:", "In which year did beijing host the Olympic Games?")

# Process query and display result
if st.button("Get Answer"):
    result = pipe(table=table, query=query)
    st.write("### Answer")
    st.write(result["answer"])