Update app.py

app.py

@@ -25,6 +25,7 @@ def get_target_style_embeddings(target_texts_batch):
 
 def run_tinystyler_batch(source_texts, target_example_texts_batch, reranking, temperature, top_p):
     inputs = tokenizer(source_texts, return_tensors="pt")
+    target_style_embeddings = get_target_style_embeddings(target_texts_batch)
     
     # Generate the output with specified temperature and top_p
     output = model.generate(
@@ -38,14 +39,15 @@ def run_tinystyler_batch(source_texts, target_example_texts_batch, reranking, te
     generated_texts = tokenizer.decode_batch(output, skip_special_tokens=True)
     return generated_texts
     
-def run_tinystyler(source_text, target_example_texts, reranking, temperature, top_p):
-    return run_tinystyler_batch([source_text], [target_example_texts], reranking, temperature, top_p)[0]
+def run_tinystyler(source_text, target_texts, reranking, temperature, top_p):
+    target_texts = [target_text.strip() for target_text in target_texts.split("\n")]
+    return run_tinystyler_batch([source_text], [target_texts], reranking, temperature, top_p)[0]
 
 # Preset examples with cached generations
 preset_examples = {
     "Example 1": {
         "source_text": "Once upon a time in a small village",
-        "target_example_texts": "In a land far away, there was a kingdom ruled by a wise king. Every day, the people of the kingdom would gather to listen to the king's stories, which were full of wisdom and kindness.",
+        "target_texts": "In a land far away, there was a kingdom ruled by a wise king. Every day, the people of the kingdom would gather to listen to the king's stories, which were full of wisdom and kindness.",
         "reranking": 5,
         "temperature": 1.0,
         "top_p": 1.0,
@@ -53,7 +55,7 @@ preset_examples = {
     },
     "Example 2": {
         "source_text": "The quick brown fox",
-        "target_example_texts": "A nimble, chocolate-colored fox swiftly darted through the emerald forest, weaving between trees with grace and agility.",
+        "target_texts": "A nimble, chocolate-colored fox swiftly darted through the emerald forest, weaving between trees with grace and agility.",
         "reranking": 5,
         "temperature": 0.9,
         "top_p": 0.9,
@@ -70,7 +72,7 @@ with gr.Blocks(theme="ParityError/[email protected]") as demo:
         example_dropdown = gr.Dropdown(label="Examples", choices=list(preset_examples.keys()))
     
     source_text = gr.Textbox(lines=3, placeholder="Enter the source text to transform into the target style...", label="Source Text")
-    target_example_texts = gr.Textbox(lines=5, placeholder="Enter example texts of the target style (one per line)...", label="Example Texts of the Target Style")
+    target_texts = gr.Textbox(lines=5, placeholder="Enter example texts of the target style (one per line)...", label="Example Texts of the Target Style")
     reranking = gr.Slider(1, 10, value=5, step=1, label="Re-ranking")
     temperature = gr.Slider(0.1, 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")
@@ -88,9 +90,9 @@ with gr.Blocks(theme="ParityError/[email protected]") as demo:
     )
     
     btn = gr.Button("Generate")
-    btn.click(run_tinystyler, [source_text, target_example_texts, reranking, temperature, top_p], output)
+    btn.click(run_tinystyler, [source_text, target_texts, reranking, temperature, top_p], output)
 
     # Initialize the fields with the first example
-    example_dropdown.value, (source_text.value, target_example_texts.value, reranking.value, temperature.value, top_p.value, output.value) = list(preset_examples.keys())[0], set_example(list(preset_examples.keys())[0])
+    example_dropdown.value, (source_text.value, target_texts.value, reranking.value, temperature.value, top_p.value, output.value) = list(preset_examples.keys())[0], set_example(list(preset_examples.keys())[0])
 
 demo.launch()