Less Magic Tokens

modeling_diva.py

@@ -88,17 +88,30 @@ class DiVAModel(PreTrainedModel):
             torch_dtype=torch.float16,
         )
         self.processor = AutoProcessor.from_pretrained(config_dict["reference_encoder"])
-        self.tokenizer = AutoTokenizer.from_pretrained("WillHeld/via-llama")
-        self.prefix = torch.tensor([128000, 128006, 882, 128007, 271]).to(
+        self.tokenizer = AutoTokenizer.from_pretrained(config_dict["reference_decoder"])
+        if self.tokenizer.pad_token_id == None:
+            override_token = list(self.tokenizer.added_tokens_decoder.items())[-1]
+            self.tokenizer.pad_token_id = override_token[0]
+            self.tokenizer.pad_tokn = str(override_token[1])
+        prefix, suffix = self.tokenizer.apply_chat_template(
+            [{"role": "user", "content": "PLACEHOLDER"}],
+            tokenize=False,
+            add_generation_prompt=True,
+        ).split("PLACEHOLDER")
+        non_null = [line for line in prefix.split("\n") if line.strip()]
+        prefix_tok = self.tokenizer.encode(prefix, add_special_tokens=False)
+        suffix_tok = self.tokenizer.encode(suffix, add_special_tokens=False)
+        self.prefix = torch.tensor(prefix_tok).to(
             self.llm_decoder.model.embed_tokens.weight.device
         )
 
-        self.pre_user_suffix = torch.tensor(
-            self.tokenizer.encode(
-                "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n"
-            )
+        self.pre_system = torch.tensor(
+            self.tokenizer.encode(non_null[0] + "\n", add_special_tokens=False)
+        ).to(self.llm_decoder.model.embed_tokens.weight.device)
+        self.post_system = torch.tensor(
+            self.tokenizer.encode("\n" + non_null[-1] + "\n", add_special_tokens=False)
         ).to(self.llm_decoder.model.embed_tokens.weight.device)
-        self.final_header = torch.tensor([128009, 128006, 78191, 128007, 271]).to(
+        self.final_header = torch.tensor(suffix_tok).to(
             self.llm_decoder.model.embed_tokens.weight.device
         )
         self.speech_encoder_device = speech_encoder_device
@@ -116,9 +129,7 @@ class DiVAModel(PreTrainedModel):
         **kwargs,
     ):
         if os.path.isdir(pretrained_model_name_or_path):
-            via_path = (
-                pretrained_model_name_or_path + "/model.safetensors"
-            )
+            via_path = pretrained_model_name_or_path + "/model.safetensors"
             config_path = pretrained_model_name_or_path + "/config.json"
         else:
             # Loading from huggingface repo
@@ -207,16 +218,16 @@ class DiVAModel(PreTrainedModel):
                     padding=True,
                     padding_side="right",
                 )["input_ids"],
-                device=self.pre_user_suffix.device,
+                device=self.pre_system.device,
             )
             prefix = torch.cat(
                 [
-                    self.pre_user_suffix.expand(
+                    self.pre_system.expand(
                         bsz,
                         -1,
                     ),
                     user_prompt_text,
-                    self.prefix.expand(
+                    self.post_system.expand(
                         bsz,
                         -1,
                     ),
@@ -292,11 +303,27 @@ class DiVAModel(PreTrainedModel):
 
         if text_prompt != None and text_prompt != "":
             user_prompt_text = torch.tensor(
-                self.tokenizer(text_prompt, add_special_tokens=False)["input_ids"],
-                device=self.pre_user_suffix.device,
+                self.tokenizer(
+                    text_prompt,
+                    add_special_tokens=False,
+                    padding=True,
+                    padding_side="right",
+                )["input_ids"],
+                device=self.pre_system.device,
             )
             prefix = torch.cat(
-                [self.pre_user_suffix, user_prompt_text, self.prefix], axis=0
+                [
+                    self.pre_system.expand(
+                        bsz,
+                        -1,
+                    ),
+                    user_prompt_text,
+                    self.post_system.expand(
+                        bsz,
+                        -1,
+                    ),
+                ],
+                axis=1,
             )
         else:
             prefix = self.prefix
@@ -344,14 +371,20 @@ class DiVAModel(PreTrainedModel):
                     "<|eot_id|>", ""
                 )
             else:
-                yield (self.tokenizer.decode(outs, skip_special_tokens=True).replace(
-                    "<|eot_id|>", ""
-                ), outputs)
+                yield (
+                    self.tokenizer.decode(outs, skip_special_tokens=True).replace(
+                        "<|eot_id|>", ""
+                    ),
+                    outputs,
+                )
         if not return_outputs:
             return self.tokenizer.decode(outs, skip_special_tokens=True).replace(
                 "<|eot_id|>", ""
             )
         else:
-            return (self.tokenizer.decode(outs, skip_special_tokens=True).replace(
-                "<|eot_id|>", ""
-            ), outputs)
+            return (
+                self.tokenizer.decode(outs, skip_special_tokens=True).replace(
+                    "<|eot_id|>", ""
+                ),
+                outputs,
+            )