Update model.py

model.py

@@ -56,33 +56,15 @@ tokenizer = transformers.AutoTokenizer.from_pretrained(
 #     return ''.join(texts)
 
 
-
-# def get_prompt(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> str:
-#     texts = [f'{system_prompt}\n']
-
-#     for user_input, response in chat_history[:-1]:
-#         texts.append(f'{user_input} {response}\n')
-
-#     # Getting the user input and response from the last tuple in the chat history
-#     last_user_input, last_response = chat_history[-1]
-#     texts.append(f' input: {last_user_input} {last_response} {message} response: ')
-
-#     return ''.join(texts)
-
 def get_prompt(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> str:
     texts = [f'{system_prompt}\n']
 
-    # If chat_history is not empty, process all but the last entry
-    if chat_history:
-        for user_input, response in chat_history[:-1]:
-            texts.append(f'{user_input} {response}\n')
+    for user_input, response in chat_history[:-1]:
+        texts.append(f'{user_input} {response}\n')
 
-        # Getting the user input and response from the last tuple in the chat history
-        last_user_input, last_response = chat_history[-1]
-        texts.append(f' input: {last_user_input} {last_response} {message} Response: ')
-    else:
-        # If chat_history is empty, just add the message with 'Response:' at the end
-        texts.append(f' input: {message} Response: ')
+    # Getting the user input and response from the last tuple in the chat history
+    last_user_input, last_response = chat_history[-1]
+    texts.append(f' input: {last_user_input} {last_response} {message} Response: ')
 
     return ''.join(texts)
 
@@ -99,26 +81,28 @@ def run(message: str,
         max_new_tokens: int = 256,
         temperature: float = 0.8,
         top_p: float = 0.95,
-        top_k: int = 50) -> str:
+        top_k: int = 50) -> Iterator[str]:
     prompt = get_prompt(message, chat_history, system_prompt)
     inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to(device)
 
-    # 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],
+    streamer = TextIteratorStreamer(tokenizer,
+                                    timeout=10.,
+                                    skip_prompt=True,
+                                    skip_special_tokens=True)
+    generate_kwargs = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
         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()
 
-    # 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
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        yield ''.join(outputs)