Update app.py

app.py

@@ -24,21 +24,23 @@ def predict(message, history):
         "cognitivecomputations/dolphin-2.8-mistral-7b-v02",
         trust_remote_code=True
     )
+
     model = AutoModelForCausalLM.from_pretrained(
         "cognitivecomputations/dolphin-2.8-mistral-7b-v02",
         torch_dtype="auto",
         load_in_4bit=True,
         trust_remote_code=True
     )
-    history_transformer_format = history + [[message, ""]]
 
+    history_transformer_format = history + [[message, ""]]
     system_prompt = "<|im_start|>system\nYou are discord user Jimmy Donaldson. You are an american youtuber by the name of MrBeast. Your sole purpose in life is to get people to laugh at you. You make innapropriate jokes, and aren't very funny.<|im_end|>"
+
     messages = system_prompt + "".join(["".join(["\n<|im_start|>user\n" + item[0], "<|im_end|>\n<|im_start|>assistant\n" + item[1]]) for item in history_transformer_format])
-    input_ids = tokenizer([messages], return_tensors="pt").to('cuda')
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
+
+    input_ids = tokenizer([messages], return_tensors="pt").to('cuda').input_ids
+
+    output_ids = model.generate(
         input_ids,
-        streamer=streamer,
         max_new_tokens=10000,
         do_sample=True,
         top_p=0.95,
@@ -46,16 +48,9 @@ def predict(message, history):
         temperature=0.7,
         num_beams=1
     )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    final_message = []
-    for new_token in streamer:
-        final_message.append(new_token)
-        if '<|im_end|>' in new_token:
-            break
-    
-    return ''.join(final_message)
+
+    output_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
+    return output_text
 
 @spaces.GPU(duration=120)
 def generate(content):