Remove unused model

app.py

@@ -173,130 +173,134 @@ 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):
-        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 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):
+#         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"))
+
+
+def DummyClipModel:
+    def __init__(self, text_model)
+        self.text_model = text_model
 
 def encode_text(text, model):
     text = normalize_text(text)
@@ -304,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ゼロショット）")
@@ -316,7 +320,9 @@ description_text = st.empty()
 if "model" not in st.session_state:
     description_text.text("日本語CLIPモデル読み込み中... ")
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = ClipModel("sonoisa/clip-vit-b-32-japanese-v1", device=device)
+    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)
     st.session_state.model = model
 
     print("extract dataset")
@@ -325,7 +331,8 @@ if "model" not in st.session_state:
     )
     
     print("loading dataset")
-    df = pq.read_table("clip_zeroshot_irasuto_items_20210224.parquet").to_pandas()
+    df = pq.read_table("clip_zeroshot_irasuto_items_20210224.parquet",
+        columns=["page", "description", "image_url", "image_vector"]).to_pandas()
     st.session_state.df = df
     
     # sentence_vectors = np.stack(df["sentence_vector"])