Update app.py

app.py

@@ -7,7 +7,7 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
-def generate(html, entity, website_desc, datasource, year, month, title):
+def generate(html, entity, website_desc, datasource, year, month, title, prompt):
     html_text = "html | " if html == "on" else ""
     entity_text = ""
     if entity != "":
@@ -23,12 +23,12 @@ def generate(html, entity, website_desc, datasource, year, month, title):
     month_text = "Month: " + month + " | " if month != "" else ""  
     title_text = "Title: " + title + " | " if title != "" else ""
 
-    prompt = html_text + year_text + month_text + website_desc_text + title_text + datasource_text + entity_text
+    final_prompt = html_text + year_text + month_text + website_desc_text + title_text + datasource_text + entity_text + prompt
     
     model = AutoModelForCausalLM.from_pretrained("bs-modeling-metadata/checkpoints_all_04_23", subfolder="checkpoint-30000step")
     tokenizer = AutoTokenizer.from_pretrained("bs-modeling-metadata/checkpoints_all_04_23", subfolder="tokenizer")
 
-    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = tokenizer(final_prompt, return_tensors="pt")
 
     outputs = model.generate(**inputs, max_new_tokens=128)
     return tokenizer.batch_decode(outputs, skip_special_tokens=True)
@@ -41,10 +41,11 @@ datasource = gr.Textbox(placeholder="enter a datasource", label="datasource")
 year = gr.Textbox(placeholder="enter a year", label="year")
 month = gr.Textbox(placeholder="enter a month", label="month")
 title = gr.Textbox(placeholder="enter a website title", label="website title")
+prompt = gr.Textbox(placeholder="enter a prompt", label="prompt")
 
 demo = gr.Interface(
     fn=generate,
-    inputs=[html, entity, website_desc, datasource, year, month, title],
+    inputs=[html, entity, website_desc, datasource, year, month, title, prompt],
     outputs="text",
 )
 demo.launch()