Arnaldo Mont'Alvao commited on
Commit
4e0b663
Β·
1 Parent(s): b289db3

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +47 -0
app.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from peft import PeftModel, PeftConfig
3
+ from transformers import AutoModelForCausalLM, AutoTokenizer
4
+
5
+ #Load model
6
+
7
+ peft_model_id = f"monta135/gen_ads_bloom"
8
+ config = PeftConfig.from_pretrained(peft_model_id)
9
+ model = AutoModelForCausalLM.from_pretrained(
10
+ config.base_model_name_or_path,
11
+ return_dict=True,
12
+ load_in_8bit=True,
13
+ device_map="auto",
14
+ )
15
+ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
16
+
17
+ # Load the Lora model
18
+ model = PeftModel.from_pretrained(model, peft_model_id)
19
+
20
+
21
+ def make_inference(product_name, product_description):
22
+ batch = tokenizer(
23
+ f"### Product and Description:\n{product_name}: {product_description}\n\n### Ad:",
24
+ return_tensors="pt",
25
+ )
26
+
27
+ with torch.cuda.amp.autocast():
28
+ output_tokens = model.generate(**batch, max_new_tokens=50)
29
+
30
+ return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
31
+
32
+
33
+ if __name__ == "__main__":
34
+
35
+ # gradio interface
36
+ import gradio as gr
37
+
38
+ gr.Interface(
39
+ make_inference,
40
+ [
41
+ gr.inputs.Textbox(lines=2, label="Product Name"),
42
+ gr.inputs.Textbox(lines=5, label="Product Description"),
43
+ ],
44
+ gr.outputs.Textbox(label="Ad"),
45
+ title="GenAds-AI",
46
+ description="GenAds-AI is a generative model that creates ads for home products.",
47
+ ).launch()