Update app/tapas.py

app/tapas.py

@@ -5,66 +5,64 @@ import re
 p = re.compile('\d+(\.\d+)?')
 
 def load_model_and_tokenizer():
-    """
-    Load
-    """
-    tokenizer = AutoTokenizer.from_pretrained("Meena/table-question-answering-tapas")
-    model = AutoModelForTableQuestionAnswering.from_pretrained("Meena/table-question-answering-tapas")
+  """
+  Load
+  """
+  tokenizer = AutoTokenizer.from_pretrained("Meena/table-question-answering-tapas")
+  model = AutoModelForTableQuestionAnswering.from_pretrained("Meena/table-question-answering-tapas")
 
-    # 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.
-    """
-    table = table.astype('str').head(100)
-    inputs = tokenizer(table=table, queries=queries, padding='max_length', return_tensors="pt")
-    return table, inputs
+  """
+  Convert dictionary into data frame and tokenize inputs given queries.
+  """
+  table = table.astype('str').head(100)
+  inputs = tokenizer(table=table, queries=queries, padding='max_length', return_tensors="pt")
+  return table, inputs
 
 
 def generate_predictions(inputs, model, tokenizer):
-    """
-    Generate predictions for some tokenized input.
-    """
-    # 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(
-          inputs,
-          outputs.logits.detach(),
-          outputs.logits_aggregation.detach()
-    )
+  """
+  Generate predictions for some tokenized input.
+  """
+  # Generate model results
+  outputs = model(**inputs)
 
-    # Return values
-    return predicted_table_cell_coords, predicted_aggregation_operators
+  # 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
 
 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):
-        if len(coordinates) == 1:
-          # 1 cell
-          answers.append(table.iat[coordinates[0]])
-        else:
-          # > 1 cell
-          cell_values = []
-          for coordinate in coordinates:
-            cell_values.append(table.iat[coordinate])
-          answers.append(", ".join(cell_values))
-          
-    # Return values
-    return aggregation_predictions_string, answers
+  """
+  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):
+      if len(coordinates) == 1:
+        # 1 cell
+        answers.append(table.iat[coordinates[0]])
+      else:
+        # > 1 cell
+        cell_values = []
+        for coordinate in coordinates:
+          cell_values.append(table.iat[coordinate])
+        answers.append(", ".join(cell_values))
+         
+  # Return values
+  return aggregation_predictions_string, answers
 
 
 def show_answers(queries, answers, aggregation_predictions_string):
@@ -90,12 +88,12 @@ def show_answers(queries, answers, aggregation_predictions_string):
   return results
 
 def execute_query(query, table):
-	"""
+  """
 	  Invoke the TAPAS model.
 	"""
-    queries = [query]
-    tokenizer, model = load_model_and_tokenizer()
-	table, inputs = prepare_inputs(table, queries, tokenizer)
-	predicted_table_cell_coords, predicted_aggregation_operators = generate_predictions(inputs, model, tokenizer)
-	aggregation_predictions_string, answers = postprocess_predictions(predicted_aggregation_operators, predicted_table_cell_coords, table)
-	return show_answers(queries, answers, aggregation_predictions_string)
+  queries = [query]
+  tokenizer, model = load_model_and_tokenizer()
+  table, inputs = prepare_inputs(table, queries, tokenizer)
+  predicted_table_cell_coords, predicted_aggregation_operators = generate_predictions(inputs, model, tokenizer)
+  aggregation_predictions_string, answers = postprocess_predictions(predicted_aggregation_operators, predicted_table_cell_coords, table)
+  return show_answers(queries, answers, aggregation_predictions_string)