Added multi-vector embeddings

1 files changed, 3 insertions, 3 deletions

diff --git a/models/clip/prompt.py b/models/clip/prompt.py
index da33ecf..9da3955 100644
--- a/models/clip/prompt.py
+++ b/models/clip/prompt.py
@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import Union
 
 import torch
 
@@ -10,13 +10,13 @@ class PromptProcessor():
         self.tokenizer = tokenizer
         self.text_encoder = text_encoder
 
-    def get_input_ids(self, prompt: Union[str, List[str]]):
+    def get_input_ids(self, prompt: Union[str, list[str]]):
         return self.tokenizer(
             prompt,
             padding="do_not_pad",
         ).input_ids
 
-    def unify_input_ids(self, input_ids: List[int]):
+    def unify_input_ids(self, input_ids: list[int]):
         return self.tokenizer.pad(
             {"input_ids": input_ids},
             padding=True,