🎨 refactor: update type alias to TextLLMModel and add COLORS_LLM list

LlmKira · Sep 26, 2024 · f4af2e6 · f4af2e6
1 parent c2d266e
commit f4af2e6
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/src/novelai_python/_enum.py b/src/novelai_python/_enum.py
@@ -62,21 +62,32 @@ class TextTokenizerGroup(object):
     TextLLMModel.ERATO: TextTokenizerGroup.LLAMA3,
 }
 
+COLORS_LLM = [
+    TextLLMModel.BLUE,
+    TextLLMModel.RED,
+    TextLLMModel.GREEN,
+    TextLLMModel.PURPLE,
+    TextLLMModel.PINK,
+    TextLLMModel.YELLOW,
+    TextLLMModel.WHITE,
+    TextLLMModel.BLACK,
+]
 
-def get_llm_group(model: TextLLMModelTypeAlias) -> Optional[TextTokenizerGroup]:
+
+def get_llm_group(model: TextLLMModel) -> Optional[TextTokenizerGroup]:
     if isinstance(model, str):
         model = TextLLMModel(model)
     return TOKENIZER_MODEL_MAP.get(model, None)
 
 
-def get_tokenizer_model(model: TextLLMModelTypeAlias) -> str:
+def get_tokenizer_model(model: TextLLMModel) -> str:
     if isinstance(model, str):
         model = TextLLMModel(model)
     group = TOKENIZER_MODEL_MAP.get(model, TextTokenizerGroup.GPT2)
     return group
 
 
-def get_tokenizer_model_url(model: TextLLMModelTypeAlias) -> str:
+def get_tokenizer_model_url(model: TextLLMModel) -> str:
     model_name = get_tokenizer_model(model)
     if not model_name.endswith(".def"):
         model_name = f"{model_name}.def"