File size: 1,510 Bytes
1ce01d1
 
 
dc07482
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1ce01d1
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
from transformers import TapasTokenizer, TFTapasForQuestionAnswering
import pandas as pd

def execute_query(query, table):
	if table:
		#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