refactor text generation code

app.py

@@ -5,7 +5,8 @@ import gradio as gr
 
 import torch
 import transformers
-from transformers import AutoTokenizer, AutoModelForCausalLM
+# from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline
 from transformers import pipeline
 from diffusers import StableDiffusionPipeline
 
@@ -29,30 +30,20 @@ def safety_checker(images, clip_input):
 pipe.safety_checker = safety_checker
 
 SAVED_CHECKPOINT = 'mikegarts/distilgpt2-lotr'
-model = AutoModelForCausalLM.from_pretrained(SAVED_CHECKPOINT)
-tokenizer = AutoTokenizer.from_pretrained(SAVED_CHECKPOINT)
+text_generation_pipe = pipeline("text-generation", model=SAVED_CHECKPOINT)
 
 summarizer = pipeline("summarization")
 
 #######################################################
 
+#######################################################
+
 def break_until_dot(txt):
     return txt.rsplit('.', 1)[0] + '.'
 
 def generate(prompt):
-    input_context = prompt
-    input_ids = tokenizer.encode(input_context, return_tensors="pt").to(model.device)
-
-    outputs = model.generate(
-        input_ids=input_ids, 
-        max_length=120,
-        min_length=50,
-        temperature=0.7, 
-        num_return_sequences=3, 
-        do_sample=True
-    )
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return break_until_dot(decoded)
+    generated = text_generation_pipe(prompt, max_length=140)[0]['generated_text']
+    return break_until_dot(generated)
 
 
 def generate_story(prompt):
@@ -110,8 +101,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.'
-                              f' The text2img model is {model_id}. The summarization is done using 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 the default distilbart.')
     prompt = gr.Textbox(label="Your prompt", value="Frodo took the sword and")
     story = gr.Textbox(label="Your story")
     summary = gr.Textbox(label="Summary")
@@ -122,7 +113,7 @@ with gr.Blocks() as demo:
     img_description = gr.Markdown('Image generation takes some time'
                                   ' but here you can see what is generated from the latent state of the diffuser every few steps.'
                                   ' Usually there is a significant improvement around step 12 that yields a much better image')
-    status_msg = gr.Markdown()
+    status_msg = gr.Markdown('')
     
     gallery = gr.Gallery()
     image = gr.Image(label='Illustration for your story', show_label=True)
@@ -147,4 +138,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)