Update app.py

app.py

@@ -5,25 +5,38 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 peft_model_id = f"mazenlhm/MarketingModel"
 config = PeftConfig.from_pretrained(peft_model_id)
-model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='auto')
+model = AutoModelForCausalLM.from_pretrained(
+    config.base_model_name_or_path,
+    return_dict=True,
+    load_in_8bit=True,
+    device_map="auto",
+)
 tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
 # Load the Lora model
 model = PeftModel.from_pretrained(model, peft_model_id)
 
-rom IPython.display import display, Markdown
 
-def generate_text(inputs):
-  product_name, description = inputs  
-  batch = tokenizer(f"Below is a product and description, please write a marketing email for this product.\n\n### Product:\n{product}\n### Description:\n{description}\n\n### Marketing Email:\n", return_tensors='pt')
+def make_inference(product_name, product_description):
+    batch = tokenizer(
+        f"### Product and Description:\n{product_name}: {product_description}\n\n### Ad:",
+        return_tensors="pt",
+    )
 
-  with torch.cuda.amp.autocast():
-    output_tokens = model.generate(**batch, max_new_tokens=200)
+    with torch.cuda.amp.autocast():
+        output_tokens = model.generate(**batch, max_new_tokens=50)
 
-  generated_text = Markdown((tokenizer.decode(output_tokens[0], skip_special_tokens=True)))
-  return generated_text
-    
-inputs = gr.inputs.Text(label="Product Name"), gr.inputs.Text(label="Description")
-output = gr.outputs.Text(label="Generated Text")
-iface = gr.Interface(fn=generate_text, inputs=inputs, outputs=output, title="Product Text Generation")
-iface.launch()
+    return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+
+
+
+gr.Interface(
+    make_inference,
+        [
+            gr.inputs.Textbox(lines=2, label="Product Name"),
+            gr.inputs.Textbox(lines=5, label="Product Description"),
+        ],
+        gr.outputs.Textbox(label="Ad"),
+        title="GenerAd-AI",
+        description="GenerAd-AI is a generative model that generates ads for products.",
+).launch()