Update all

app.py

@@ -14,7 +14,7 @@ import pyminizip
 import transformers
 from transformers import BertJapaneseTokenizer, BertModel
 from huggingface_hub import hf_hub_download, snapshot_download
-# from PIL import Image
+from PIL import Image
 
 
 def unicode_normalize(cls, s):
@@ -172,135 +172,135 @@ class ClipTextModel(nn.Module):
         torch.save(self.output_linear.state_dict(), os.path.join(output_dir, "output_linear.bin"))
 
 
-# class ClipVisionModel(nn.Module):
-#     def __init__(self, model_name_or_path, device=None):
-#         super(ClipVisionModel, self).__init__()
-
-#         if os.path.exists(model_name_or_path):
-#             # load from file system
-#             visual_projection_state_dict = torch.load(os.path.join(model_name_or_path, "visual_projection.bin"))
-#         else:
-#             # download from the Hugging Face model hub
-#             filename = hf_hub_download(repo_id=model_name_or_path, filename="visual_projection.bin")
-#             visual_projection_state_dict = torch.load(filename)
-
-#         self.model = transformers.CLIPVisionModel.from_pretrained(model_name_or_path)
-#         config = self.model.config
-
-#         self.feature_extractor = transformers.CLIPFeatureExtractor.from_pretrained(model_name_or_path)
-
-#         vision_embed_dim = config.hidden_size
-#         projection_dim = 512
-
-#         self.visual_projection = nn.Linear(vision_embed_dim, projection_dim, bias=False)
-#         self.visual_projection.load_state_dict(visual_projection_state_dict)
-
-#         self.eval()
-
-#         if device is None:
-#             device = "cuda" if torch.cuda.is_available() else "cpu"
-#         self.device = torch.device(device)
-#         self.to(self.device)
-
-#     def forward(
-#         self,
-#         pixel_values=None,
-#         output_attentions=None,
-#         output_hidden_states=None,
-#         return_dict=None,
-#     ):
-#         output_states = self.model(
-#             pixel_values=pixel_values,
-#             output_attentions=output_attentions,
-#             output_hidden_states=output_hidden_states,
-#             return_dict=return_dict,
-#         )
-#         image_embeds = self.visual_projection(output_states[1])
-
-#         return image_embeds
-
-#     @torch.no_grad()
-#     def encode_image(self, images, batch_size=8):
-#         self.eval()
-#         all_embeddings = []
-#         iterator = range(0, len(images), batch_size)
-#         for batch_idx in iterator:
-#             batch = images[batch_idx:batch_idx + batch_size]
-
-#             encoded_input = self.feature_extractor(batch, return_tensors="pt").to(self.device)
-#             model_output = self(**encoded_input)
-#             image_embeddings = model_output.cpu()
-
-#             all_embeddings.extend(image_embeddings)
-
-#         # return torch.stack(all_embeddings).numpy()
-#         return torch.stack(all_embeddings)        
-
-#     @staticmethod
-#     def remove_alpha_channel(image):
-#         image.convert("RGBA")
-#         alpha = image.convert('RGBA').split()[-1]
-#         background = Image.new("RGBA", image.size, (255, 255, 255))
-#         background.paste(image, mask=alpha)
-#         image = background.convert("RGB")
-#         return image
-
-#     def save(self, output_dir):
-#         self.model.save_pretrained(output_dir)
-#         self.feature_extractor.save_pretrained(output_dir)
-#         torch.save(self.visual_projection.state_dict(), os.path.join(output_dir, "visual_projection.bin"))
-
-
-# class ClipModel(nn.Module):
-#     def __init__(self, model_name_or_path, device=None):
-#         super(ClipModel, self).__init__()
-
-#         if os.path.exists(model_name_or_path):
-#             # load from file system
-#             repo_dir = model_name_or_path
-#         else:
-#             # download from the Hugging Face model hub
-#             repo_dir = snapshot_download(model_name_or_path)
-
-#         self.text_model = ClipTextModel(repo_dir, device=device)
-#         self.vision_model = ClipVisionModel(os.path.join(repo_dir, "vision_model"), device=device)
-
-#         with torch.no_grad():
-#             logit_scale = nn.Parameter(torch.ones([]) * 2.6592)
-#             logit_scale.set_(torch.load(os.path.join(repo_dir, "logit_scale.bin")).clone().cpu())
-#             self.logit_scale = logit_scale
-
-#         self.eval()
-
-#         if device is None:
-#             device = "cuda" if torch.cuda.is_available() else "cpu"
-#         self.device = torch.device(device)
-#         self.to(self.device)
-
-#     def forward(self, pixel_values, input_ids, attention_mask, token_type_ids):
-#         image_features = self.vision_model(pixel_values=pixel_values)
-#         text_features = self.text_model(input_ids=input_ids, 
-#                                         attention_mask=attention_mask, 
-#                                         token_type_ids=token_type_ids)[0]
-
-#         image_features = image_features / image_features.norm(dim=-1, keepdim=True)
-#         text_features = text_features / text_features.norm(dim=-1, keepdim=True)
-
-#         logit_scale = self.logit_scale.exp()
-#         logits_per_image = logit_scale * image_features @ text_features.t()
-#         logits_per_text = logits_per_image.t()
-
-#         return logits_per_image, logits_per_text
-
-#     def save(self, output_dir):
-#         torch.save(self.logit_scale, os.path.join(output_dir, "logit_scale.bin"))
-#         self.text_model.save(output_dir)
-#         self.vision_model.save(os.path.join(output_dir, "vision_model"))
-
-
-class DummyClipModel:
-    def __init__(self, text_model):
-        self.text_model = text_model
+class ClipVisionModel(nn.Module):
+    def __init__(self, model_name_or_path, device=None):
+        super(ClipVisionModel, self).__init__()
+
+        if os.path.exists(model_name_or_path):
+            # load from file system
+            visual_projection_state_dict = torch.load(os.path.join(model_name_or_path, "visual_projection.bin"))
+        else:
+            # download from the Hugging Face model hub
+            filename = hf_hub_download(repo_id=model_name_or_path, filename="visual_projection.bin")
+            visual_projection_state_dict = torch.load(filename)
+
+        self.model = transformers.CLIPVisionModel.from_pretrained(model_name_or_path)
+        config = self.model.config
+
+        self.feature_extractor = transformers.CLIPFeatureExtractor.from_pretrained(model_name_or_path)
+
+        vision_embed_dim = config.hidden_size
+        projection_dim = 512
+
+        self.visual_projection = nn.Linear(vision_embed_dim, projection_dim, bias=False)
+        self.visual_projection.load_state_dict(visual_projection_state_dict)
+
+        self.eval()
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = torch.device(device)
+        self.to(self.device)
+
+    def forward(
+        self,
+        pixel_values=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        output_states = self.model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        image_embeds = self.visual_projection(output_states[1])
+
+        return image_embeds
+
+    @torch.no_grad()
+    def encode_image(self, images, batch_size=8):
+        self.eval()
+        all_embeddings = []
+        iterator = range(0, len(images), batch_size)
+        for batch_idx in iterator:
+            batch = images[batch_idx:batch_idx + batch_size]
+
+            encoded_input = self.feature_extractor(batch, return_tensors="pt").to(self.device)
+            model_output = self(**encoded_input)
+            image_embeddings = model_output.cpu()
+
+            all_embeddings.extend(image_embeddings)
+
+        # return torch.stack(all_embeddings).numpy()
+        return torch.stack(all_embeddings)        
+
+    @staticmethod
+    def remove_alpha_channel(image):
+        image.convert("RGBA")
+        alpha = image.convert('RGBA').split()[-1]
+        background = Image.new("RGBA", image.size, (255, 255, 255))
+        background.paste(image, mask=alpha)
+        image = background.convert("RGB")
+        return image
+
+    def save(self, output_dir):
+        self.model.save_pretrained(output_dir)
+        self.feature_extractor.save_pretrained(output_dir)
+        torch.save(self.visual_projection.state_dict(), os.path.join(output_dir, "visual_projection.bin"))
+
+
+class ClipModel(nn.Module):
+    def __init__(self, model_name_or_path, device=None):
+        super(ClipModel, self).__init__()
+
+        if os.path.exists(model_name_or_path):
+            # load from file system
+            repo_dir = model_name_or_path
+        else:
+            # download from the Hugging Face model hub
+            repo_dir = snapshot_download(model_name_or_path)
+
+        self.text_model = ClipTextModel(repo_dir, device=device)
+        self.vision_model = ClipVisionModel(os.path.join(repo_dir, "vision_model"), device=device)
+
+        with torch.no_grad():
+            logit_scale = nn.Parameter(torch.ones([]) * 2.6592)
+            logit_scale.set_(torch.load(os.path.join(repo_dir, "logit_scale.bin")).clone().cpu())
+            self.logit_scale = logit_scale
+
+        self.eval()
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = torch.device(device)
+        self.to(self.device)
+
+    def forward(self, pixel_values, input_ids, attention_mask, token_type_ids):
+        image_features = self.vision_model(pixel_values=pixel_values)
+        text_features = self.text_model(input_ids=input_ids, 
+                                        attention_mask=attention_mask, 
+                                        token_type_ids=token_type_ids)[0]
+
+        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
+
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        return logits_per_image, logits_per_text
+
+    def save(self, output_dir):
+        torch.save(self.logit_scale, os.path.join(output_dir, "logit_scale.bin"))
+        self.text_model.save(output_dir)
+        self.vision_model.save(os.path.join(output_dir, "vision_model"))
+
+
+# class DummyClipModel:
+#     def __init__(self, text_model):
+#         self.text_model = text_model
 
 def encode_text(text, model):
     text = normalize_text(text)
@@ -308,10 +308,10 @@ def encode_text(text, model):
     return text_embedding
 
 
-# def encode_image(image_filename, model):
-#     image = Image.open(image_filename)
-#     image_embedding = model.vision_model.encode_image([image]).numpy()
-#     return image_embedding
+def encode_image(image_filename, model):
+    image = Image.open(image_filename)
+    image_embedding = model.vision_model.encode_image([image]).numpy()
+    return image_embedding
 
 
 st.title("いらすと検索（日本語CLIPゼロショット）")
@@ -321,30 +321,31 @@ if "model" not in st.session_state:
     description_text.text("日本語CLIPモデル読み込み中... ")
     device = "cuda" if torch.cuda.is_available() else "cpu"
     text_model = ClipTextModel("sonoisa/clip-vit-b-32-japanese-v1", device=device)
-    # model = ClipModel("sonoisa/clip-vit-b-32-japanese-v1", device=device)
-    model = DummyClipModel(text_model)
+    model = ClipModel("sonoisa/clip-vit-b-32-japanese-v1", device=device)
+    # model = DummyClipModel(text_model)
     st.session_state.model = model
 
     print("extract dataset")
     pyminizip.uncompress(
-        "clip_zeroshot_irasuto_image_items_20210224.pq.zip", st.secrets["ZIP_PASSWORD"], None, 1
+        "clip_zeroshot_irasuto_items_20210224.pq.zip", st.secrets["ZIP_PASSWORD"], None, 1
     )
     
     print("loading dataset")
-    df = pq.read_table("clip_zeroshot_irasuto_image_items_20210224.parquet",
-        columns=["page", "description", "image_url", "image_vector"]).to_pandas()
-    st.session_state.df = df
+    df = pq.read_table("clip_zeroshot_irasuto_items_20210224.parquet",
+        columns=["page", "description", "image_url", "sentence_vector", "image_vector"]).to_pandas()
     
-    # sentence_vectors = np.stack(df["sentence_vector"])
+    sentence_vectors = np.stack(df["sentence_vector"])
     image_vectors = np.stack(df["image_vector"])
-    # st.session_state.sentence_vectors = sentence_vectors
+    st.session_state.sentence_vectors = sentence_vectors
+
+    st.session_state.df = df
     st.session_state.image_vectors = image_vectors
 
     print("finished loading model and dataset")
     
 model = st.session_state.model
 df = st.session_state.df
-# sentence_vectors = st.session_state.sentence_vectors
+sentence_vectors = st.session_state.sentence_vectors
 image_vectors = st.session_state.image_vectors
 
 description_text.text("日本語CLIPモデル（ゼロショット）を用いて、説明文の意味が近い「いらすとや」画像を検索します。\nキーワードを列挙するよりも、自然な文章を入力した方が精度よく検索できます。\n画像は必ずリンク先の「いらすとや」さんのページを開き、そこからダウンロードしてください。")

clip_zeroshot_irasuto_image_items_20210224.pq.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2f602399369a485f1586b7ca04e8ae096868ecce85527928671b08bf5e80c200
-size 54262882
+oid sha256:62eabf2fd3664a3ddfe29bb7ee59027fa37a34a1d05a9704f09ac363ad5acb2f
+size 72554784

clip_zeroshot_irasuto_items_20210224.pq.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:321d909ba0f92425a5107ad26a6d97dc4f7601b2b4f22ab020199f2ba2237ce7
-size 104296063
+oid sha256:3059351ecc86353c53ba25f7cb5e74db0e55b1ba5257402970a20fd04158b5f1
+size 122826331

requirements.txt

@@ -4,4 +4,4 @@ pyminizip
 fugashi
 ipadic
 scipy
-#pillow==7.1.2
+pillow==7.1.2