feat(generate): Make prepending BOS model-conigurable

And disable it for Granite Code models

Branch: GraniteCodeSupport

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

Author: gabe-l-hart

torchchat/generate.py

@@ -746,6 +746,7 @@ def encode_tokens(self, string, bos=True, device="cpu"):
         if bos:
             tokens = [self.tokenizer.bos_id()] + tokens
         logger.debug("Size after encode_tokens: %d", len(tokens))
+        logger.debug("Token IDs: %s", tokens)
         return torch.tensor(tokens, dtype=torch.int, device=device)
 
     def _callback(self, x, *, buffer, done_generating):
@@ -794,7 +795,7 @@ def _gen_model_input(
             # Single String prompt
             if isinstance(prompt, str):
                 encoded = self.encode_tokens(
-                    prompt, bos=True, device=self.builder_args.device
+                    prompt, bos=self.model.config.tokenizer_prepend_bos, device=self.builder_args.device
                 )
             # List of dialog
             else:
@@ -1048,7 +1049,7 @@ def chat(
                     else:
                         prompt = f"{B_INST} {prompt.strip()} {E_INST}"
                     encoded = self.encode_tokens(
-                        prompt, bos=True, device=self.builder_args.device
+                        prompt, bos=self.model.config.tokenizer_prepend_bos, device=self.builder_args.device
                     )
                 else:
                     if self.system_prompt:

torchchat/model.py

@@ -276,6 +276,7 @@ class TransformerArgs:
     # Select the desired tokenizer. Defaults to sentencepiece
     use_tiktoken: bool = False
     use_hf_tokenizer: bool = False
+    tokenizer_prepend_bos: bool = True
     max_seq_length: int = 8192
     rope_scaling: Optional[Dict[str, Any]] = None
     # For pipeline parallel
@@ -333,13 +334,15 @@ class ModelArgs:
     transformer_args: Dict[str, Dict[str, Any]]
     use_tiktoken: bool
     use_hf_tokenizer: bool
+    tokenizer_prepend_bos: bool
 
     def __init__(
         self,
         transformer_args: Dict[str, Dict[str, Any]],
         model_type: ModelType = ModelType.TextOnly,
         use_tiktoken: bool = False,
         use_hf_tokenizer: bool = False,
+        tokenizer_prepend_bos: bool = True,
     ) -> None:
         self._sanity_check(transformer_args, model_type)
 
@@ -349,6 +352,7 @@ def __init__(
         # Model-level attributes
         self.use_tiktoken = use_tiktoken
         self.use_hf_tokenizer = use_hf_tokenizer
+        self.tokenizer_prepend_bos = tokenizer_prepend_bos
 
     def _sanity_check(
         self,
@@ -376,7 +380,14 @@ def from_params(cls, params_path):
 
         use_tiktoken = loaded_params.get("use_tiktoken", False)
         use_hf_tokenizer = loaded_params.get("use_hf_tokenizer", False)
-        return cls(transformer_args, model_type, use_tiktoken, use_hf_tokenizer)
+        tokenizer_prepend_bos = loaded_params.get("tokenizer_prepend_bos", True)
+        return cls(
+            transformer_args=transformer_args,
+            model_type=model_type,
+            use_tiktoken=use_tiktoken,
+            use_hf_tokenizer=use_hf_tokenizer,
+            tokenizer_prepend_bos=tokenizer_prepend_bos,
+        )
 
     @classmethod
     def from_table(cls, name: str):

torchchat/model_params/Granite-3B-Code.json

@@ -8,6 +8,7 @@
     "rope_base": 10000000,
     "vocab_size": 49152,
     "use_hf_tokenizer": true,
+    "tokenizer_prepend_bos": false,
     "norm_eps": 0.00001,
     "rope_scaling": null,
     "attention_bias": true,

torchchat/model_params/Granite-8B-Code.json

@@ -8,6 +8,7 @@
     "rope_base": 10000000,
     "vocab_size": 49152,
     "use_hf_tokenizer": true,
+    "tokenizer_prepend_bos": false,
     "norm_eps": 0.00001,
     "rope_scaling": null,
     "attention_bias": true,