Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
-from threading import Thread
 
 
 title = "🦅Falcon 🗨️ChatBot"
@@ -9,7 +8,7 @@ description = "Falcon-RW-1B is a 1B parameters causal decoder-only model built b
 examples = [["How are you?"]]
 
 
-tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-rw-1b",torch_dtype=torch.float16)
+tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-rw-1b")
 model = AutoModelForCausalLM.from_pretrained(
     "tiiuae/falcon-rw-1b",
     trust_remote_code=True,
@@ -30,13 +29,16 @@ def predict(message, history):
     history_transformer_format = history + [[message, ""]]
     stop = StopOnTokens()
 
+    #Construct the input message string for the model by concatenating the current system message and conversation history
     messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])  #curr_system_message + 
                 for item in history_transformer_format])
     
+    #Tokenize the messages string
     model_inputs = tokenizer([messages], return_tensors="pt")
     streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         model_inputs,
+        streamer=streamer,
         max_new_tokens=1024,
         do_sample=True,
         top_p=0.95,
@@ -45,9 +47,11 @@ def predict(message, history):
         num_beams=1,
         stopping_criteria=StoppingCriteriaList([stop])
         )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
+    #t = Thread(target=model.generate, kwargs=generate_kwargs)
+    #t.start()
+    model.generate(**generate_kwargs)
 
+    #Initialize an empty string to store the generated text
     partial_message  = ""
     for new_token in streamer:
         if new_token != '<':
@@ -64,4 +68,4 @@ gr.ChatInterface(predict,
     undo_btn="Delete Previous",
     clear_btn="Clear",
     chatbot=gr.Chatbot(height=300),
-    textbox=gr.Textbox(placeholder="Chat with me")).launch()
+    textbox=gr.Textbox(placeholder="Chat with me").queue().launch()