File size: 894 Bytes
6d20a1a
4c83d7c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from transformers import TapasTokenizer, TapexTokenizer, BartForConditionalGeneration
import pandas as pd
import datetime


def execute_query(query, csv_file):
    a = datetime.datetime.now()

    table = pd.read_csv(csv_file.name, delimiter=",")
    table.fillna(0, inplace=True)
    table = table.astype(str)

    model_name = "microsoft/tapex-large-finetuned-wtq"
    model = BartForConditionalGeneration.from_pretrained(model_name)
    tokenizer = TapexTokenizer.from_pretrained(model_name)

    queries = [query]

    encoding = tokenizer(table=table, queries=queries, padding=True, return_tensors="tf",truncated=True)
    outputs = model.generate(**encoding)
    ans = tokenizer.batch_decode(outputs, skip_special_tokens=True)



    query_result = {
        "query": query,
        "answer": ans[0]
    }

    b = datetime.datetime.now()
    print(b - a)

    return query_result, table