Cleanup

1 files changed, 4 insertions, 1 deletions

diff --git a/models/clip/tokenizer.py b/models/clip/tokenizer.py
index 39c41ed..789b525 100644
--- a/models/clip/tokenizer.py
+++ b/models/clip/tokenizer.py
@@ -55,6 +55,9 @@ def shuffle_auto(tokens: list[int]):
     return shuffle_all(tokens)
 
 
+ShuffleAlgorithm = Union[bool, Literal["all", "trailing", "leading", "between", "off"]]
+
+
 class MultiCLIPTokenizer(CLIPTokenizer):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -73,7 +76,7 @@ class MultiCLIPTokenizer(CLIPTokenizer):
     def set_dropout(self, dropout: float):
         self.dropout = dropout
 
-    def set_use_vector_shuffle(self, algorithm: Union[bool, Literal["all", "trailing", "leading", "between", "off"]]):
+    def set_use_vector_shuffle(self, algorithm: ShuffleAlgorithm):
         if algorithm == "leading":
             self.vector_shuffle = shuffle_leading
         elif algorithm == "trailing":