update

README.md

@@ -7,7 +7,6 @@ tags:
 - clip
 - japanese-clip
 ---
-
 # recruit-jp/japanese-clip-vit-b-32-roberta-base
 
 ## Overview
@@ -41,17 +40,19 @@ pip install pillow requests transformers torch torchvision sentencepiece
 ```python
 import io
 import requests
-from PIL import Image
 
 import torch
 import torchvision
+from PIL import Image
 from transformers import AutoTokenizer, AutoModel
 
+
 model_name = "recruit-jp/japanese-clip-vit-b-32-roberta-base"
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModel.from_pretrained(model_name, trust_remote_code=True).to(device)
 
+
 def _convert_to_rgb(image):
     return image.convert('RGB')
 
@@ -68,25 +69,29 @@ preprocess = torchvision.transforms.Compose([
 def tokenize(tokenizer, texts):
     texts = ["[CLS]" + text for text in texts]
     encodings = [
+        # NOTE: the maximum token length that can be fed into this model is 77
         tokenizer(text, max_length=77, padding="max_length", truncation=True, add_special_tokens=False)["input_ids"]
         for text in texts
     ]
     return torch.LongTensor(encodings)
 
+
 # Run!
-image = Image.open(io.BytesIO(requests.get('https://images.pexels.com/photos/2253275/pexels-photo-2253275.jpeg?auto=compress&cs=tinysrgb&dpr=3&h=750&w=1260').content))
+image = Image.open(
+    io.BytesIO(
+        requests.get(
+            'https://images.pexels.com/photos/2253275/pexels-photo-2253275.jpeg?auto=compress&cs=tinysrgb&dpr=3&h=750&w=1260'
+        ).content
+    )
+)
 image = preprocess(image).unsqueeze(0).to(device)
 text = tokenize(tokenizer, texts=["犬", "猫", "象"]).to(device)
-
 with torch.inference_mode():
     image_features = model.get_image_features(image)
     image_features /= image_features.norm(dim=-1, keepdim=True)
-
     text_features = model.get_text_features(input_ids=text)
     text_features /= text_features.norm(dim=-1, keepdim=True)
-
     probs = image_features @ text_features.T
-
 print("Label probs:", probs.cpu().numpy()[0])
 ```
 

modeling_japanese_clip.py

@@ -84,7 +84,9 @@ class AttentionalPooler(nn.Module):
     ):
         super().__init__()
         self.query = nn.Parameter(torch.randn(n_queries, d_model))
-        self.attn = nn.MultiheadAttention(d_model, n_head, kdim=context_dim, vdim=context_dim)
+        self.attn = nn.MultiheadAttention(
+            d_model, n_head, kdim=context_dim, vdim=context_dim
+        )
         self.ln_q = norm_layer(d_model)
         self.ln_k = norm_layer(context_dim)
 
@@ -92,7 +94,9 @@ class AttentionalPooler(nn.Module):
         x = self.ln_k(x).permute(1, 0, 2)  # NLD -> LND
         N = x.shape[1]
         q = self.ln_q(self.query)
-        out = self.attn(q.unsqueeze(1).expand(-1, N, -1), x, x, need_weights=False)[0]
+        out = self.attn(
+            q.unsqueeze(1).expand(-1, N, -1), x, x, need_weights=False
+        )[0]
         return out.permute(1, 0, 2)  # LND -> NLD
 
 
@@ -187,7 +191,12 @@ class Transformer(nn.Module):
 
         self.resblocks = nn.ModuleList([
             ResidualAttentionBlock(
-                width, heads, mlp_ratio, ls_init_value=ls_init_value, act_layer=act_layer, norm_layer=norm_layer)
+                width, 
+                heads, 
+                mlp_ratio, 
+                ls_init_value=ls_init_value, 
+                act_layer=act_layer, 
+                norm_layer=norm_layer)
             for _ in range(layers)
         ])