Add EOS token interruption

app.py

@@ -44,6 +44,10 @@ def generate_and_visualize(prompt, num_tokens=10):
 
         input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)
         input_embeds = model.get_input_embeddings()(input_ids)
+
+        if next_token.item() == tokenizer.eos_token_id:
+            print("EOS token generated, stopping generation.")
+            break
     input_tokens = really_clean_tokens(tokenizer.convert_ids_to_tokens(input_ids[0]))
     generated_tokens = really_clean_tokens(tokenizer.convert_ids_to_tokens(generated_tokens_ids))
     
@@ -198,10 +202,10 @@ How can you build tools simply in transformers? Just use the decorator""",
 ]
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Token Generation with Hover Notes")
+    gr.Markdown("# Syntax highlighted text generation - for RAG applications")
     
     input_text = gr.Textbox(label="Enter your prompt:", lines=10, value=examples[0])
-    num_tokens = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Number of tokens to generate")
+    num_tokens = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Number of tokens to generate (while no EOS token)")
     generate_button = gr.Button("Generate")
 
     output_html = gr.HTML(label="Generated Output")