Update app.py

app.py

@@ -14,9 +14,10 @@ prompter_model, prompter_tokenizer = load_prompter()
 def generate(plain_text):
     input_ids = prompter_tokenizer(plain_text.strip()+" Rephrase:", return_tensors="pt").input_ids
     eos_id = prompter_tokenizer.eos_token_id
-    outputs = prompter_model.generate(input_ids, do_sample=False, max_new_tokens=75, num_beams=8, num_return_sequences=8, eos_token_id=eos_id, pad_token_id=eos_id, length_penalty=-1.0)
-    output_texts = prompter_tokenizer.batch_decode(outputs, skip_special_tokens=True)
-    res = output_texts[0].replace(plain_text+" Rephrase:", "").strip()
+    # Just use 1 beam and get 1 output, this is much, much, much faster than 8 beams and 8 outputs and we're only using the first.
+    outputs = prompter_model.generate(input_ids, do_sample=False, max_new_tokens=75, eos_token_id=eos_id, pad_token_id=eos_id, length_penalty=-1.0)
+    # Use [input_ids.shape[-1]:] because the decoded tokenised version of plain_text may have a different number of characters to the original
+    res = tokenizer.decode(outputs[0][input_ids.shape[-1]:])
     return res
 
 txt = grad.Textbox(lines=1, label="Initial Text", placeholder="Input Prompt")