adding text streaming

app.py

@@ -1,7 +1,7 @@
 import os
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline, TextStreamer
 
 token = os.environ["HUGGINGFACEHUB_API_TOKEN"]
 
@@ -51,15 +51,17 @@ def get_prompt_with_template(message: str) -> str:
 def generate_model_response(message: str) -> str:
     prompt = get_prompt_with_template(message)
     inputs = tokenizer(prompt, return_tensors='pt')
+    streamer = TextStreamer(tokenizer)
     if torch.cuda.is_available():
         inputs = inputs.to('cuda')
     # Include **generate_kwargs to include the user-defined options
     output = model.generate(**inputs, 
                             max_new_tokens=4096, 
                             do_sample=True,
-                            temperature=0.1
+                            temperature=0.1,
+                            streamer=streamer
                             ) 
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    # return tokenizer.decode(output[0], skip_special_tokens=True)
 
 def extract_response_content(full_response: str) -> str:
     response_start_index = full_response.find("### Assistant:")