Team Ai
Apppublic

IThinkUPC/demo_sqlgenerator

sourceHugging Faceupdated 3y agoView on Hugging Face
0likes
app.py26 linesDownload Raw Back to root
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()