Update app.py

app.py

@@ -51,14 +51,10 @@ def convert_history(chat_history, max_input_length=1024):
 @torch.inference_mode()
 def instruct(instruction, max_token_output=1024):
     input_text = instruction
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
-    input_ids = tokenizer(input_text, return_tensors='pt', truncation=False)
-    input_ids["input_ids"] = input_ids["input_ids"].cuda()
-    input_ids["attention_mask"] = input_ids["attention_mask"].cuda()
-    generation_kwargs = dict(input_ids, streamer=streamer, max_new_tokens=max_token_output, do_sample=False)
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    return streamer
+    input_ids = tokenizer.encode(input_text, return_tensors='pt').cuda()
+    output_ids = model.generate(input_ids, max_length=max_token_output)[0]
+    output_str = tokenizer.decode(output_ids[input_ids.shape[-1]:])
+    return output_str.strip()
 
 
 with gr.Blocks() as demo: