Update app/tapas.py

app/tapas.py

@@ -0,0 +1,48 @@
+from transformers import TapasTokenizer, TFTapasForQuestionAnswering
+import pandas as pd
+
+def execute_query(query, csv_file):
+
+    table = pd.read_csv(csv_file.name, delimiter=",")
+    table.fillna(0, inplace=True)
+    table = table.astype(str)
+
+    model_name = "google/tapas-base-finetuned-wtq"
+    model = TFTapasForQuestionAnswering.from_pretrained(model_name)
+    tokenizer = TapasTokenizer.from_pretrained(model_name)
+
+    queries = [query]
+
+    inputs = tokenizer(table=table, queries=queries, padding="max_length", return_tensors="tf")
+    outputs = model(**inputs)
+
+    predicted_answer_coordinates, predicted_aggregation_indices = tokenizer.convert_logits_to_predictions(
+        inputs, outputs.logits, outputs.logits_aggregation
+    )
+
+    # let's print out the results:
+    id2aggregation = {0: "NONE", 1: "SUM", 2: "AVERAGE", 3: "COUNT"}
+    aggregation_predictions_string = [id2aggregation[x] for x in predicted_aggregation_indices]
+
+    answers = []
+    for coordinates in predicted_answer_coordinates:
+        if len(coordinates) == 1:
+            # only a single cell:
+            answers.append(table.iat[coordinates[0]])
+        else:
+            # multiple cells
+            cell_values = []
+            for coordinate in coordinates:
+                cell_values.append(table.iat[coordinate])
+            answers.append(cell_values)
+
+    for query, answer, predicted_agg in zip(queries, answers, aggregation_predictions_string):
+        if predicted_agg != "NONE":
+            answers.append(predicted_agg)
+
+    query_result = {
+        "query": query,
+        "result": answers
+    }
+
+    return query_result, table