text encoder entirely set through the path provided in the config

Author: harshg99

src/vseek/tools/server.py

@@ -49,6 +49,7 @@ def __init__(self, args):
         if self.retriever_name == "ViClip":
             self.retriever = ViClip(
                 pretrained_model_path=self.retrieval_model_path,
+                tokenizer_path=args.retriever.text_encoder_model_path,
                 gpu_number=self.gpu_number,
             )
         else:

src/vseek/video_embedding/lib_viclip/__init__.py

@@ -7,9 +7,10 @@
 
 
 def get_viclip(size='l', 
-               pretrain=os.path.join(os.path.dirname(os.path.abspath(__file__)), "ViClip-InternVid-10M-FLT.pth")):
+               pretrain=os.path.join(os.path.dirname(os.path.abspath(__file__)), "ViClip-InternVid-10M-FLT.pth"),
+               tokenizer_path=None):
 
-    tokenizer = _Tokenizer()
+    tokenizer = _Tokenizer(bpe_path=tokenizer_path) if tokenizer_path else _Tokenizer()
     vclip = ViCLIP(tokenizer=tokenizer, size=size, pretrain=pretrain)
     m = {'viclip':vclip, 'tokenizer':tokenizer}
 

src/vseek/video_embedding/lib_viclip/viclip.py

@@ -217,13 +217,15 @@ def build_text_encoder(self):
                 context_length=self.max_txt_l,
                 vocab_size=self.text_encoder_vocab_size,
                 checkpoint_num=0,
+                tokenizer=self.tokenizer,
             )
         elif encoder_name == "vit_b16":
             text_encoder = clip_text_b16(
                 pretrained=self.text_encoder_pretrained,
                 context_length=self.max_txt_l,
                 vocab_size=self.text_encoder_vocab_size,
                 checkpoint_num=0,
+                tokenizer=self.tokenizer,
             )
         else:
             raise NotImplementedError(f"Not implemented: {encoder_name}")

src/vseek/video_embedding/lib_viclip/viclip_text.py

@@ -107,11 +107,12 @@ def __init__(
         transformer_heads: int,
         transformer_layers: int,
         checkpoint_num: int,
+        tokenizer=None,
     ):
         super().__init__()
 
         self.context_length = context_length
-        self._tokenizer = _Tokenizer()
+        self._tokenizer = tokenizer if tokenizer else _Tokenizer()
 
         self.transformer = Transformer(
             width=transformer_width,
@@ -211,6 +212,7 @@ def clip_text_b16(
     transformer_layers=12,
     checkpoint_num=0,
     pretrained=True,
+    tokenizer=None,
 ):
     # raise NotImplementedError
     model = CLIP_TEXT(
@@ -221,6 +223,7 @@ def clip_text_b16(
         transformer_heads,
         transformer_layers,
         checkpoint_num,
+        tokenizer=tokenizer,
     )
     # pretrained = _MODELS["ViT-B/16"]
     # logger.info(f"Load pretrained weights from {pretrained}")
@@ -269,6 +272,7 @@ def clip_text_l14(
     transformer_layers=12,
     checkpoint_num=0,
     pretrained=True,
+    tokenizer=None,
 ):
     model = CLIP_TEXT(
         embed_dim,
@@ -278,6 +282,7 @@ def clip_text_l14(
         transformer_heads,
         transformer_layers,
         checkpoint_num,
+        tokenizer=tokenizer,
     )
     if pretrained:
         if isinstance(pretrained, str) and pretrained != "bert-base-uncased":

src/vseek/video_embedding/video_clip.py

@@ -27,13 +27,14 @@ def __init__(
         pretrained_model_path: str,
         size: str = "l",
         gpu_number: int = 0,
+        tokenizer_path: str = None,
     ) -> None:
         """Initialize the ViClip model."""
         self.model_cfg = {
             "size": size,
             "pretrained": pretrained_model_path,
         }
-        self.model = get_viclip(self.model_cfg["size"], self.model_cfg["pretrained"])
+        self.model = get_viclip(self.model_cfg["size"], self.model_cfg["pretrained"], tokenizer_path=tokenizer_path)
         self.device = get_device(gpu_number)
         self.clip = self.model["viclip"]
         self.tokenizer = self.model["tokenizer"]