|
from transformers import AutoModelWithLMHead, AutoTokenizer |
|
|
|
tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL") |
|
model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL") |
|
|
|
def get_sql(query): |
|
input_text = "translate English to SQL: %s </s>" % query |
|
features = tokenizer([input_text], return_tensors='pt') |
|
|
|
output = model.generate(input_ids=features['input_ids'], |
|
attention_mask=features['attention_mask']) |
|
|
|
return tokenizer.decode(output[0]) |
|
|
|
query = "How many models were finetuned using BERT as base model?" |
|
|
|
res = get_sql(query) |
|
|
|
print(res) |
|
|
|
|
|
|
