Update app.py

app.py

@@ -14,7 +14,7 @@ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
 # Load the Lora model
 model = PeftModel.from_pretrained(model, peft_model_id)
-
+'''
 def make_inference(product, description):
   batch = tokenizer(f"### INSTRUCTION\nBelow 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')
 
@@ -23,19 +23,16 @@ def make_inference(product, description):
   return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
 '''
 def make_inference(product_name, product_description):
-    batch = tokenizer(
-        f"### Product and Description:\n{product_name}: {product_description}\n\n### Ad:",
-        return_tensors="pt",
-    )
-    
+    batch = tokenizer(f"### INSTRUCTION\nBelow 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')
+
     batch = {key: value.to('cuda:0') for key, value in batch.items()}
 
     with torch.cuda.amp.autocast():
-        output_tokens = model.generate(**batch, max_new_tokens=50)
+        output_tokens = model.generate(**batch, max_new_tokens=200)
 
     return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
 
-'''
+
 
 # Load the Lora model
 model = PeftModel.from_pretrained(model, peft_model_id)