Update app/tapas.py

app/tapas.py

@@ -7,69 +7,64 @@ import re
 
 p = re.compile('\d+(\.\d+)?')
 
-# Define the questions
-queries = [
-            "When did Spider-Man: No Way Home	 release?",
-            "which Movies have rating 5?"
-           ]
 
 def load_model_and_tokenizer():
-  """
+    """
     Load
-  """
-  # Load pretrained tokenizer: TAPAS finetuned on WikiTable Questions
-  tokenizer = TapasTokenizer.from_pretrained("google/tapas-base-finetuned-wtq")
+    """
+    # Load pretrained tokenizer: TAPAS finetuned on WikiTable Questions
+    tokenizer = TapasTokenizer.from_pretrained("google/tapas-base-finetuned-wtq")
 
-  # Load pretrained model: TAPAS finetuned on WikiTable Questions
-  model = TapasForQuestionAnswering.from_pretrained("google/tapas-base-finetuned-wtq")
+    # Load pretrained model: TAPAS finetuned on WikiTable Questions
+    model = TapasForQuestionAnswering.from_pretrained("google/tapas-base-finetuned-wtq")
 
-  # Return tokenizer and model
-  return tokenizer, model
+    # Return tokenizer and model
+    return tokenizer, model
 
 
 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
+    """
+    # 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
 
 
 def generate_predictions(inputs, model, tokenizer):
-  """
+    """
     Generate predictions for some tokenized input.
-  """
-  # Generate model results
-  outputs = model(**inputs)
+    """
+    # Generate model results
+    outputs = model(**inputs)
 
-  # Convert logit outputs into predictions for table cells and aggregation operators
-  predicted_table_cell_coords, predicted_aggregation_operators = tokenizer.convert_logits_to_predictions(
+    # Convert logit outputs into predictions for table cells and aggregation operators
+    predicted_table_cell_coords, predicted_aggregation_operators = tokenizer.convert_logits_to_predictions(
           inputs,
           outputs.logits.detach(),
           outputs.logits_aggregation.detach()
-  )
-  
-  # Return values
-  return predicted_table_cell_coords, predicted_aggregation_operators
+    )
+
+    # Return values
+    return predicted_table_cell_coords, predicted_aggregation_operators
 
 
 def postprocess_predictions(predicted_aggregation_operators, predicted_table_cell_coords, table):
-  """
+    """
     Compute the predicted operation and nicely structure the answers.
-  """
-  # Process predicted aggregation operators
-  aggregation_operators = {0: "NONE", 1: "SUM", 2: "AVERAGE", 3:"COUNT"}
-  aggregation_predictions_string = [aggregation_operators[x] for x in predicted_aggregation_operators]
-
-  # Process predicted table cell coordinates
-  answers = []
-  for agg, coordinates in zip(predicted_aggregation_operators, predicted_table_cell_coords):
+    """
+    # Process predicted aggregation operators
+    aggregation_operators = {0: "NONE", 1: "SUM", 2: "AVERAGE", 3:"COUNT"}
+    aggregation_predictions_string = [aggregation_operators[x] for x in predicted_aggregation_operators]
+
+    # Process predicted table cell coordinates
+    answers = []
+    for agg, coordinates in zip(predicted_aggregation_operators, predicted_table_cell_coords):
     if len(coordinates) == 1:
       # 1 cell
       answers.append(table.iat[coordinates[0]])
@@ -80,30 +75,30 @@ def postprocess_predictions(predicted_aggregation_operators, predicted_table_cel
         cell_values.append(table.iat[coordinate])
       answers.append(", ".join(cell_values))
       
-  # Return values
-  return aggregation_predictions_string, answers
+    # Return values
+    return aggregation_predictions_string, answers
 
 
 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)}
+    agg = {"NONE": lambda x: x, "SUM" : lambda x: sum(x), "AVERAGE": lambda x: (sum(x) / len(x)), "COUNT": lambda x: len(x)}
+
+    for query, answer, predicted_agg in zip(queries, answers, aggregation_predictions_string):
+        print(query)
+        if predicted_agg == "NONE":
+            print("Predicted answer: " + answer)
+        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(','))))))
+                return "Predicted answer: " + 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(','))))
+                return "Predicted answer: " + str(agg[predicted_agg](answer.split(',')))
+            else:
+                return "Predicted answer: " + predicted_agg + " > " + answer
 
-  for query, answer, predicted_agg in zip(queries, answers, aggregation_predictions_string):
-    print(query)
-    if predicted_agg == "NONE":
-      print("Predicted answer: " + answer)
-    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(','))))))
-        return "Predicted answer: " + 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(','))))
-        return "Predicted answer: " + str(agg[predicted_agg](answer.split(',')))
-      else:
-		return "Predicted answer: " + predicted_agg + " > " + answer
-