Update model.py

model.py

@@ -68,30 +68,26 @@ def run(message: str,
         max_new_tokens: int = 256,
         temperature: float = 0.8,
         top_p: float = 0.95,
-        top_k: int = 50) -> Iterator[str]:
+        top_k: int = 50) -> str:
     prompt = get_prompt(message, chat_history, system_prompt)
     inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to(device)
 
-    streamer = TextIteratorStreamer(tokenizer,
-                                    timeout=10.,
-                                    skip_prompt=True,
-                                    skip_special_tokens=True)
-    generate_kwargs = dict(
-        inputs,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
+    # Generate tokens using the model
+    output = model.generate(
+        input_ids=inputs['input_ids'],
+        attention_mask=inputs['attention_mask'],
+        max_length=max_new_tokens + inputs['input_ids'].shape[-1],
         do_sample=True,
         top_p=top_p,
         top_k=top_k,
         temperature=temperature,
-        num_beams=1,
+        num_beams=1
     )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        if "instruction:" in text:
-            break
-        yield ''.join(outputs)
+
+    # Decode the output tokens back to a string
+    output_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    
+    # Remove everything including and after "instruct: "
+    output_text = output_text.split("instruct: ")[0]
+
+    return output_text