gotta catch em all

viclip.py

@@ -11,20 +11,37 @@ from .simple_tokenizer import SimpleTokenizer as _Tokenizer
 from .viclip_vision import clip_joint_l14, clip_joint_b16
 from .viclip_text import clip_text_l14, clip_text_b16
 
-logger = logging.getLogger(__name__)
+# from transformers import AutoModel
+from transformers import PreTrainedModel #new
+from transformers import PretrainedConfig
 
+logger = logging.getLogger(__name__)
 
-class ViCLIP(nn.Module):
-    """docstring for ViCLIP"""
+from .configuration_viclip import Config
+# class ViCLIP(nn.Module):
+class ViCLIP(PreTrainedModel):
+    _auto_class="AutoModel"
+    config_class=Config
+
+    def __init__(self,
+                #  tokenizer=None,  # config:PretrainedConfig is the only parameter
+                #  size='l',
+                #  pretrain=None,
+                #  freeze_text=True,
+                 config=PretrainedConfig()):
+        super(ViCLIP, self).__init__(config)
+        self.config=config
+        if 'size' in config.to_dict(): ###########
+            size=config.size
+            pretrain=None
+            tokenizer_path=config.tokenizer_path
+            tokenizer=None
+            freeze_text=True
 
-    def __init__(self,  
-                 tokenizer=None, 
-                 size='l',
-                 pretrain=os.path.join(os.path.dirname(os.path.abspath(__file__)), "ViClip-InternVid-10M-FLT.pth"),
-                 freeze_text=True):
-        super(ViCLIP, self).__init__()
         if tokenizer:
             self.tokenizer = tokenizer
+        elif tokenizer_path:
+            self.tokenizer = _Tokenizer(tokenizer_path)
         else:
             self.tokenizer = _Tokenizer()
         self.max_txt_l = 32
@@ -217,6 +234,7 @@ class ViCLIP(nn.Module):
                 context_length=self.max_txt_l,
                 vocab_size=self.text_encoder_vocab_size,
                 checkpoint_num=0,
+                tokenizer_path=None if not 'tokenizer_path' in self.config.to_dict() else self.config.tokenizer_path
             )
         elif encoder_name == "vit_b16":
             text_encoder = clip_text_b16(
@@ -224,6 +242,7 @@ class ViCLIP(nn.Module):
                 context_length=self.max_txt_l,
                 vocab_size=self.text_encoder_vocab_size,
                 checkpoint_num=0,
+                tokenizer_path=None if not 'tokenizer_path' in self.config.to_dict() else self.config.tokenizer_path
             )
         else:
             raise NotImplementedError(f"Not implemented: {encoder_name}")
@@ -253,10 +272,10 @@ class ViCLIP(nn.Module):
         return clip_feat
 
     def get_predict_label(self, clip_feature, text_feats_tensor, top=5):
-        label_probs = (100.0 * clip_feature @ text_feats_tensor.T)
+        label_probs = (100.0 * clip_feature @ text_feats_tensor.T).softmax(dim=-1)
         top_probs, top_labels = label_probs.cpu().topk(top, dim=-1)
         return top_probs, top_labels
 
     
 if __name__ =="__main__":
-    tokenizer = _Tokenizer()
+    tokenizer = _Tokenizer()