feat: add adapter model with self attention

data_search/adapter_utils.py

@@ -2,18 +2,30 @@ import torch
 import torch.nn as nn
 
 
+class Adapter(nn.Module):
+    def __init__(self, img_dim, txt_dim, embed_dim=1024, num_heads=8):
+        super().__init__()
+        self.adapter = nn.Sequential(
+            nn.Linear(img_dim, embed_dim),
+            nn.ReLU(),
+            nn.Linear(embed_dim, embed_dim),
+            nn.ReLU(),
+            nn.Linear(embed_dim, txt_dim)
+        )
+        self.self_attention = nn.MultiheadAttention(embed_dim=txt_dim, num_heads=num_heads)
+
+    def forward(self, img_emb):
+        img_emb = self.adapter(img_emb).unsqueeze(0)
+        attn_output, _= self.self_attention(img_emb, img_emb, img_emb)
+        return attn_output.squeeze(0)
+
+
 def get_adapter_model(in_shape, out_shape):
-    model = nn.Sequential(
-        nn.Linear(in_shape, 1024),
-        nn.ReLU(),
-        nn.Linear(1024, 1024),
-        nn.ReLU(),
-        nn.Linear(1024, out_shape)
-    )
+    model = Adapter(in_shape, out_shape)
     return model
 
 
 def load_adapter_model():
     model = get_adapter_model(512, 384)
-    model.load_state_dict(torch.load("./weights/adapter_model.pt", map_location=torch.device('cpu')))
+    model.load_state_dict(torch.load("./weights/adapter_model_with_attention.pt", map_location=torch.device('cpu')))
     return model

data_upload/data_upload_page.py

@@ -9,6 +9,7 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
 def data_upload(clip_model, preprocess, text_embedding_model):
     st.title("Data Upload")
+    st.warning("Please note that this is a public application. Make sure you are not uploading any sensitive data.")
     upload_choice = st.selectbox(options=["Upload Image", "Add Image from URL / Link", "Upload PDF", "Website Link"], label="Select Upload Type")
     if upload_choice == "Upload Image":
         image_util.upload_image(clip_model, preprocess)

weights/adapter_model_with_attention.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32785c4b2a93e21c358e57465021938c9c85e5cfc4bf477aff33afc72ad62cb2
+size 10243596