Update Flask app on each step of generated text from neural network model

I have a trained PyTorch model that can generate text given a text input. To be more specific, users can insert a title and the model will generate the body for it to form an article. Right now, I can only make my Flask app render when the entire body is generated (which takes quite a long time), instead of having it update on each step of the generation process. What I am aiming for is what can be found at Talk To Transformer, where the model output is periodically generated until completion.

I only know very basic web skills with ajax, php, and jQuery, and I was wondering how I can achieve this through Flask and minimal backend technologies.

So I've tried the following, but obviously it doesn't work because it only generates HTML when render_template is followed by a return statement, but I think it shows what I'm trying to do. Even if I could get the page to render at each prediction step of the model, I feel like it's a hacky way to do it so I'm willing to learn about how this could be done in a professional/best way.

@app.route('/main', methods=['GET', 'POST'])
def main(): 
  # <some code for preparing the model>

  # when the user input is passed as a POST request to this page: 
  if request.method =='POST': 
    print(request.form)
    title = request.form.get('title', None)
    if title:
      src = tokenizer.encode(title) 
      current_output = [] 
      with torch.no_grad():
        while len(current_output) < MAX_LEN: 
          next_token = predict_next(model, tokenizer, src, current_output, args)
          if next_token:
            current_output.append(next_token)
            output = tokenizer.decode(current_output, skip_special_tokens=True)
            render_template('index.html', title=title, output=output)
          else: 
            break

      return render_template('index.html', title=title, output=output)

if __name__ == '__main__':
  app.run(host='0.0.0.0', port=403, debug=True)

Any pointers are greatly appreciated!!