Update app.py

app.py

@@ -6,7 +6,6 @@ from threading import Thread
 # Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")
 model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")
-model = model.to('cuda:0')
 
 class StopOnTokens(StoppingCriteria):
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
@@ -24,7 +23,7 @@ def predict(message, history):
     messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])  #curr_system_message +
                 for item in history_transformer_format])
 
-    model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
+    model_inputs = tokenizer([messages], return_tensors="pt")
     temperature = gr.Slider(minimum=0.1, maximum=2.0, value=1.0, label="Temperature"),
     max_new_tokens = gr.Slider(minimum=0, maximum=2048, value=10, label="Temperature"),
     min_new_tokens = gr.Slider(minimum=0, maximum=2048, value=1, label="Temperature"),
@@ -39,14 +38,10 @@ def predict(message, history):
         num_beams=1,
         stopping_criteria=StoppingCriteriaList([stop])
         )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    partial_message  = ""
-    for new_token in streamer:
-        if new_token != '<':
-            partial_message += new_token
-            yield partial_message
+    generated_sequence = model.generate(**generate_kwargs)[0]
+    generated_text = tokenizer.decode(generated_sequence, skip_special_tokens=True)
+
+    yield generated_text
 
 
 gr.ChatInterface(predict).queue().launch()