Update app.py

app.py

@@ -6,7 +6,7 @@ model_name = "google/flan-t5-large"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-def concatenate_and_generate(text1, text2, temperature, top_p):
+def concatenate_and_generate(text1, text2, reranking, temperature, top_p):
     concatenated_text = text1 + " " + text2
     inputs = tokenizer(concatenated_text, return_tensors="pt")
     
@@ -27,14 +27,15 @@ with gr.Blocks(theme="ParityError/[email protected]") as demo:
     gr.Markdown("# TinyStyler Demo")
     gr.Markdown("Style transfer the source text into the target style, given some example texts of the target style. You can adjust re-ranking and top_p to your desire to control the quality of style transfer. A higher re-ranking value will generally result in better results, at slower speed.")
 
-    text1 = gr.Textbox(lines=2, placeholder="Enter the source text to transform into the target style...")
-    text2 = gr.Textbox(lines=2, placeholder="Enter example texts of the target style (one per line)...")
-    temperature = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature")
-    top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.1, label="Top-p")
+    text1 = gr.Textbox(lines=3, placeholder="Enter the source text to transform into the target style...", label="Source Text")
+    text2 = gr.Textbox(lines=5, placeholder="Enter example texts of the target style (one per line)...", label="Examples of Target Style")
+    reranking = gr.Slider(1, 10, value=5, step=1, label="Re-ranking")
+    temperature = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="Temperature")
+    top_p = gr.Slider(0.0, 1.0, value=1.0, step=0.1, label="Top-P")
     
     output = gr.Textbox()
     
     btn = gr.Button("Generate")
-    btn.click(concatenate_and_generate, [text1, text2, temperature, top_p], output)
+    btn.click(concatenate_and_generate, [text1, text2, reranking, temperature, top_p], output)
 
 demo.launch()