fix(chat): Fix small formatting bugs in llama3 chat formatter

Branch: GraniteCodeSupport

Signed-off-by: Gabe Goodhart <[email protected]>

Author: gabe-l-hart

torchchat/generate.py

@@ -110,7 +110,7 @@ def _encode_message(self, message: _ChatFormatter.MESSAGE_TYPE) -> List[int]:
                         self.tokenizer.encode(content["text"], bos=False, eos=False)
                     )
 
-        tokens.append(self.tokenizer.special_tokens["<|eot_id|>"])
+        tokens.append(self.tokenizer.special_tokens["<|eot_id|>\n"])
         return tokens
 
     def encode_dialog_prompt(
@@ -123,8 +123,8 @@ def encode_dialog_prompt(
         for message in dialog:
             tokens.extend(self._encode_message(message))
         # Add the start of an assistant message for the model to complete.
-        if add_generation_prompt:
-            tokens.extend(self._encode_header("assistant"))  # Pass role directly as a string
+        if add_generation_prompt and dialog and dialog[-1]["role"] != "assistant":
+            tokens.extend(self._encode_header("assistant")) # Pass role directly as a string
         return tokens