Update app.py

app.py

@@ -3,6 +3,7 @@ import pandas as pd
 import torch
 from transformers import pipeline
 from transformers import TapasTokenizer, TapasForQuestionAnswering
+import datetime
 
 df = pd.read_excel('discrepantes.xlsx', index_col='Unnamed: 0')
 df.fillna(0, inplace=True)
@@ -10,17 +11,27 @@ table_data = df.astype(str)
 print(table_data.head())
 
 def response(user_question, table_data):
-    tokenizer = TapasTokenizer.from_pretrained("google/tapas-large-finetuned-wtq", drop_rows_to_fit=True)
-    model = TapasForQuestionAnswering.from_pretrained("google/tapas-large-finetuned-wtq")
+    a = datetime.datetime.now()
 
-    inputs = tokenizer(table=table_data, queries=user_question, padding="max_length", truncation=True, return_tensors="pt")
-    outputs = model(**inputs)
+    model_name = "microsoft/tapex-large-finetuned-wtq"
+    model = BartForConditionalGeneration.from_pretrained(model_name)
+    tokenizer = TapexTokenizer.from_pretrained(model_name)
 
-    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]]
+    queries = [user_question]
 
-    return predicted_answer_coordinates, aggregation_predictions
+    encoding = tokenizer(table=table_data, query=queries, padding=True, return_tensors="pt",truncation=True)
+    outputs = model.generate(**encoding)
+    ans = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+    query_result = {
+        "query": query,
+        "answer": ans[0]
+    }
+
+    b = datetime.datetime.now()
+    print(b - a)
+
+    return query_result, table
 
 # Streamlit interface
 st.markdown("""