monicacherto commited on
Commit
7c364cd
Β·
1 Parent(s): f52b54a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -8
app.py CHANGED
@@ -2,17 +2,24 @@ import gradio as gr
2
  import torch
3
  from peft import PeftModel, PeftConfig
4
 
 
 
 
 
 
 
 
 
5
  def greet(name):
6
  return "Hello " + name + "!!"
7
 
8
- def model_run():
9
- peft_model_id = f"IThinkUPC/SQLGenerator-AI"
10
- config = PeftConfig.from_pretrained(peft_model_id)
11
- model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='auto')
12
- tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
13
- # Load the Lora model
14
- model = PeftModel.from_pretrained(model, peft_model_id)
15
 
16
 
17
- iface = gr.Interface(fn=greet, inputs="text", outputs="text")
 
18
  iface.launch()
 
2
  import torch
3
  from peft import PeftModel, PeftConfig
4
 
5
+
6
+ peft_model_id = f"IThinkUPC/SQLGenerator-AI"
7
+ config = PeftConfig.from_pretrained(peft_model_id)
8
+ model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='auto')
9
+ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
10
+ # Load the Lora model
11
+ model = PeftModel.from_pretrained(model, peft_model_id)
12
+
13
  def greet(name):
14
  return "Hello " + name + "!!"
15
 
16
+ def make_inference(prompt):
17
+ batch = tokenizer(f"### Question:\n{prompt}: \n\n### Query", return_tensors='pt')
18
+ with torch.cuda.amp.autocast():
19
+ output_tokens = model.generate(**batch, max_new_tokens=50)
20
+ dreturn tokenizer.decode(output_tokens[0], skip_special_tokens=True)
 
 
21
 
22
 
23
+ #iface = gr.Interface(fn=greet, inputs="text", outputs="text")
24
+ iface = gr.Interface(fn=make_inference, inputs="text", outputs="text")
25
  iface.launch()