Update app/tapas.py

app/tapas.py

@@ -9,12 +9,7 @@ def load_model_and_tokenizer():
     """
     Load
     """
-    # Load pretrained tokenizer: TAPAS finetuned on WikiTable Questions
-    # tokenizer = TapasTokenizer.from_pretrained("google/tapas-base-finetuned-wtq")
     tokenizer = AutoTokenizer.from_pretrained("Meena/table-question-answering-tapas")
-
-    # Load pretrained model: TAPAS finetuned on WikiTable Questions
-    # model = TapasForQuestionAnswering.from_pretrained("google/tapas-base-finetuned-wtq")
     model = AutoModelForTableQuestionAnswering.from_pretrained("Meena/table-question-answering-tapas")
 
     # Return tokenizer and model
@@ -25,13 +20,8 @@ def prepare_inputs(table, queries, tokenizer):
     """
     Convert dictionary into data frame and tokenize inputs given queries.
     """
-    # Prepare inputs
-    # table = pd.DataFrame.from_dict(data)
-    # table = netflix_df[['title', 'release_year', 'rating']].astype('str').head(50)
     table = table.astype('str').head(100)
     inputs = tokenizer(table=table, queries=queries, padding='max_length', return_tensors="pt")
-
-    # Return things
     return table, inputs
 
 
@@ -83,7 +73,7 @@ def show_answers(queries, answers, aggregation_predictions_string):
     Visualize the postprocessed answers.
   """
   agg = {"NONE": lambda x: x, "SUM" : lambda x: sum(x), "AVERAGE": lambda x: (sum(x) / len(x)), "COUNT": lambda x: len(x)}
-  result = ''
+  results = []
   for query, answer, predicted_agg in zip(queries, answers, aggregation_predictions_string):
       print(query)
       if predicted_agg == "NONE":
@@ -91,13 +81,14 @@ def show_answers(queries, answers, aggregation_predictions_string):
       else:
           if all([not p.match(val) == None for val in answer.split(', ')]):
               # print("Predicted answer: " + predicted_agg + "(" + answer + ") = " + str(agg[predicted_agg](list(map(float, answer.split(','))))))
-              result = "Predicted answer: " + str(agg[predicted_agg](list(map(float, answer.split(',')))))
+              result = str(agg[predicted_agg](list(map(float, answer.split(',')))))
           elif predicted_agg == "COUNT":
               # print("Predicted answer: " + predicted_agg + "(" + answer + ") = " + str(agg[predicted_agg](answer.split(','))))
-              result = "Predicted answer: " + str(agg[predicted_agg](answer.split(',')))
+              result = str(agg[predicted_agg](answer.split(',')))
           else:
-              result = "Predicted answer: " + predicted_agg + " > " + answer
-  return result
+              result = predicted_agg + " > " + answer
+          results.append(result)
+  return results
 
 def execute_query(queries, table):