IThinkUPC/demo_sqlgenerator
0
1import gradio as gr2import torch3from peft import PeftModel, PeftConfig4from transformers import AutoModelForCausalLM, AutoTokenizer5 6 7peft_model_id = f"IThinkUPC/SQLGenerator-AI"8config = PeftConfig.from_pretrained(peft_model_id)9model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='auto')10tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)11# Load the Lora model12model = PeftModel.from_pretrained(model, peft_model_id)13 14def greet(name):15 return "Hello " + name + "!!"16 17def make_inference(prompt): 18 batch = tokenizer(f"### Question:\n{prompt}: \n\n### Query", return_tensors='pt')19 with torch.cuda.amp.autocast():20 output_tokens = model.generate(**batch, max_new_tokens=50)21 return tokenizer.decode(output_tokens[0], skip_special_tokens=True)22 23 24#iface = gr.Interface(fn=greet, inputs="text", outputs="text")25iface = gr.Interface(fn=make_inference, inputs="text", outputs="text")26iface.launch()