Update app.py

app.py

@@ -1,6 +1,8 @@
 import streamlit as st
 import pandas as pd
+import torch
 from transformers import pipeline
+from transformers import TapasTokenizer, TapasForQuestionAnswering
 
 df = pd.read_excel('discrepantes.xlsx', index_col='Unnamed: 0')
 df.fillna(0, inplace=True)
@@ -8,9 +10,17 @@ table_data = df.astype(str)
 print(table_data.head())
 
 def response(user_question, table_data):
-    tqa = pipeline(task="table-question-answering", model="google/tapas-large-finetuned-wtq")
-    resposta = tqa(table=table_data , query=user_question)['cells'][0]
-    return resposta
+    tokenizer = TapasTokenizer.from_pretrained("google/tapas-large-finetuned-wtq", drop_rows_to_fit=True)
+    model = TapasForQuestionAnswering.from_pretrained("google/tapas-large-finetuned-wtq")
+
+    inputs = tokenizer(table=table_data, queries=user_question, padding="max_length", truncation=True, return_tensors="pt")
+    outputs = model(**inputs)
+
+    predicted_answer_coordinates = outputs.predicted_answer_coordinates.detach().cpu().numpy()
+    id2aggregation = {0: 'NONE', 1: 'SUM', 2: 'AVERAGE', 3: 'COUNT'}
+    aggregation_predictions = id2aggregation[outputs.aggregation_predictions.detach().cpu().numpy()[0]]
+
+    return predicted_answer_coordinates, aggregation_predictions
 
 # Streamlit interface
 st.markdown("""