Upload tokenization_dart.py

tokenization_dart.py

@@ -1,7 +1,5 @@
 import logging
-import json
-from typing import Dict, List
-from pydantic.dataclasses import dataclass
+from typing import List
 
 from transformers import PreTrainedTokenizerFast
 from tokenizers.decoders import Decoder
@@ -39,35 +37,24 @@ PROMPT_TEMPLATE = (
     "{{ '</character>' }}"
 
     "{{ '<general>' }}"
+    # length token
+    "{% if 'length' not in messages or messages['length'] is none %}"
+    "{{ '<|long|>' }}"
+    "{% else %}"
+    "{{ messages['length'] }}"
+    "{% endif %}"
+
+    # general token
     "{% if 'general' not in messages or messages['general'] is none %}"
     "{{ '' }}"
     "{% else %}"
     "{{ messages['general'] }}"
     "{% endif %}"
+    "{{ '<|input_end|>' }}"
 ).strip()
 # fmt: on
 
 
-@dataclass
-class Category:
-    name: str
-    bos_token_id: int
-    eos_token_id: int
-
-
-@dataclass
-class TagCategoryConfig:
-    categories: Dict[str, Category]
-    category_to_token_ids: Dict[str, List[int]]
-
-
-def load_tag_category_config(config_json: str):
-    with open(config_json, "rb") as file:
-        config: TagCategoryConfig = TagCategoryConfig(**json.loads(file.read()))
-
-    return config
-
-
 class DartDecoder:
     def __init__(self, special_tokens: List[str]):
         self.special_tokens = list(special_tokens)