Update app.py

app.py

@@ -38,18 +38,20 @@ summarizer = pipeline("summarization")
 
 #######################################################
 
-def break_until_dot(txt):
-    return txt.rsplit('.', 1)[0] + '.'
-
 def generate(prompt):
-    generated = text_generation_pipe(prompt, max_length=140)[0]['generated_text']
-    return break_until_dot(generated)
+    res = text_generation_pipe(prompt,
+               max_length=140,
+               repetition_penalty=1.1)[0]['generated_text']
 
+    i=0
+    while res[-1] != '.' and i < 30:
+        res = text_generation_pipe(res, max_length=1)[0]['generated_text']
+        i += 1
+    return res
 
 def generate_story(prompt):
     story = generate(prompt=prompt)
     summary = summarizer(story, min_length=5, max_length=15)[0]['summary_text']
-    summary = break_until_dot(summary)
     return story, summary, gr.update(visible=True)
 
 def on_change_event(app_state):
@@ -101,8 +103,8 @@ with gr.Blocks() as demo:
                          })
     title = gr.Markdown('## Lord of the rings app')
     description = gr.Markdown(f'#### A Lord of the rings inspired app that combines text and image generation.'
-                              f' The language modeling is done by fine tuning distilgpt2 on the LOTR trilogy ({SAVED_CHECKPOINT}).'
-                              f' The text2img model is {model_id}. The summarization is done using the default distilbart.')
+                              f' The language modeling is done by fine tuning distilgpt2 on the LOTR trilogy {SAVED_CHECKPOINT}.'
+                              f' The text2img model is {model_id}. The summarization is done using distilbart.')
     prompt = gr.Textbox(label="Your prompt", value="Frodo took the sword and")
     story = gr.Textbox(label="Your story")
     summary = gr.Textbox(label="Summary")
@@ -138,4 +140,4 @@ with gr.Blocks() as demo:
 if READ_TOKEN:
     demo.queue().launch()
 else:
-    demo.queue().launch(share=True, debug=True)
+    demo.queue().launch(share=True, debug=True)