app.py

import gradio as gr
from transformers import pipeline
import logging

# Set up logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

# Load the Tapas model
qa_pipeline = pipeline("table-question-answering", model="google/tapas-large-finetuned-wtq")

def ask_table(data, question):
    logging.info(f"Received Data: {data}")
    logging.info(f"Received Question: {question}")
    
    try:
        if not data or not question:
            return {"answer": "Please provide both a table and a question."}
        
        headers = data[0]  # First row is headers
        rows = data[1:]    # Remaining rows are data
        
        # Check if headers and rows exist
        if not headers or not rows:
            return {"answer": "Table data is empty or improperly formatted."}

        # Check table size (up to 50 rows, 20 columns)
        if len(rows) > 50:
            return {"answer": "The table has too many rows. Limit: 50 rows."}
        if len(headers) > 20:
            return {"answer": "The table has too many columns. Limit: 20 columns."}
        
        # Convert to TAPAS-compatible format
        processed_table = [dict(zip(headers, row)) for row in rows]
        
        if not processed_table:
            return {"answer": "Table data is empty or improperly formatted."}
        
        # Call the model to get the answer
        answers = qa_pipeline(table=processed_table, query=question)
        answer_text = answers.get("answer", "No answer found.")
        
        return {"answer": answer_text}
    
    except Exception as e:
        logging.error(f"Error processing table: {str(e)}")
        return {"answer": f"Error: {str(e)}"}

# Define Gradio interface
iface = gr.Interface(
    fn=ask_table,
    inputs=[
        gr.Dataframe(
            headers=None,  # Expect the user to provide headers in the first row
            row_count=(2, "dynamic"),
            col_count=(19, "dynamic"),
            type="array",
            label="Input Table (First row = Headers)"
        ),
        gr.Textbox(
            lines=2,
            placeholder="Enter your question about the table here...",
            label="Ask a Question"
        )
    ],
    outputs="text",
    title="Table Question Answering",
    description="Provide a table with headers in the first row and ask questions. Supports up to 50 rows and 20 columns."
)

# Launch Gradio app
iface.launch(share=True)