diff --git a/modules/tokenizer.py b/modules/tokenizer.py index 66e031d..935027d 100644 --- a/modules/tokenizer.py +++ b/modules/tokenizer.py @@ -93,9 +93,7 @@ class BpeTokenizer: @property def stop_ids(self) -> List[int]: - stop_ids = [] - for token in self._control_tokens: - stop_ids.append(self._tokenizer.token_to_id(token)) + stop_ids = self._control_tokens + self._special_tokens return stop_ids @property @@ -108,12 +106,4 @@ class BpeTokenizer: @property def pad_id(self) -> int: - return self._tokenizer.token_to_id("") - - @property - def user_id(self) -> int: - return self._tokenizer.token_to_id("<|user|>") - - @property - def system_id(self) -> int: - return self._tokenizer.token_to_id("<|system|>") \ No newline at end of file + return self._tokenizer.token_to_id("") \ No newline at end of file diff --git a/modules/utils.py b/modules/utils.py index 302a06f..2119daf 100644 --- a/modules/utils.py +++ b/modules/utils.py @@ -119,8 +119,8 @@ def get_pt_processor(tokenizer: BpeTokenizer): def get_sft_processor(tokenizer: BpeTokenizer): def processor(input_dict: dict): query, response = input_dict["query"], input_dict["response"] - q = tokenizer.encode(f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant") - a = tokenizer.encode(f"\n{response}\n") + q = tokenizer.encode(f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n") + a = tokenizer.encode(f"{response}<|im_end|>\n") tokens = torch.tensor(q + a, dtype=torch.int32) loss_mask = torch.zeros_like(tokens, dtype=torch.bool) loss_mask[len(q):] = True