from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("ansumanpandey/codgen-finetuned-SQLQueryGeneration")
model = AutoModelForCausalLM.from_pretrained("ansumanpandey/codgen-finetuned-SQLQueryGeneration")

def get_sql(query):
  input_text = "Query to %s </s>" % query
  features = tokenizer([input_text], return_tensors='pt')

  output = model.generate(input_ids=features['input_ids'],
               attention_mask=features['attention_mask'],
               max_new_tokens=70)
  sql_query= tokenizer.decode(output[0])
  return sql_query