feat: clip model test

app.py

@@ -5,25 +5,38 @@ import pandas as pd
 import streamlit as st
 import torch
 from sentence_transformers.util import semantic_search
-from transformers import VisionTextDualEncoderModel, VisionTextDualEncoderProcessor
+from transformers import AutoModel, AutoProcessor
 
 st.title("VitB32 Bert Ko Small Clip Test")
 st.markdown("Unsplash data에서 입력 텍스트와 가장 유사한 이미지를 검색합니다.")
 
 
 @st.cache(allow_output_mutation=True, show_spinner=False)
-def get_model():
+def get_dual_encoder_model():
     with st.spinner("Loading model..."):
-        model = VisionTextDualEncoderModel.from_pretrained(
-            "Bingsu/vitB32_bert_ko_small_clip"
-        ).eval()
-        processor = VisionTextDualEncoderProcessor.from_pretrained(
-            "Bingsu/vitB32_bert_ko_small_clip"
-        )
+        model = AutoModel.from_pretrained("Bingsu/vitB32_bert_ko_small_clip").eval()
+        processor = AutoProcessor.from_pretrained("Bingsu/vitB32_bert_ko_small_clip")
     return model, processor
 
 
-model, processor = get_model()
+@st.cache(allow_output_mutation=True, show_spinner=False)
+def get_clip_model():
+    with st.spinner("Loading model..."):
+        model = AutoModel.from_pretrained("Bingsu/clip-vit-base-patch32-ko").eval()
+        processor = AutoProcessor.from_pretrained("Bingsu/clip-vit-base-patch32-ko")
+    return model, processor
+
+
+model_type = st.radio(
+    "Select model",
+    ["Bingsu/clip-vit-base-patch32-ko", "Bingsu/vitB32_bert_ko_small_clip"],
+    horizontal=True,
+)
+
+if model_type == "Bingsu/clip-vit-base-patch32-ko":
+    model, processor = get_clip_model()
+else:
+    model, processor = get_dual_encoder_model()
 
 info = pd.read_csv("info.csv")
 with open("img_id.pkl", "rb") as f:
@@ -35,7 +48,6 @@ tokens = processor(text=text, return_tensors="pt")
 
 with torch.no_grad():
     text_emb = model.get_text_features(**tokens)
-text_emb = text_emb / text_emb.norm(dim=1, keepdim=True)
 
 result = semantic_search(text_emb, img_emb, top_k=15)[0]
 _result = iter(result)

pyproject.toml

@@ -1,25 +1,25 @@
 [tool.poetry]
 name = "vitb32_bert_ko_small_clip_test"
-version = "0.1.0"
+version = "0.2.0"
 description = ""
 authors = ["Bingsu <ks2515@naver.com>"]
 license = "MIT"
 
 [tool.poetry.dependencies]
 python = "^3.9"
-torch = "^1.11.0"
-transformers = "^4.19.2"
-streamlit = "^1.10.0"
-pandas = "^1.4.2"
-sentence-transformers = "^2.2.0"
+torch = "^1.12"
+transformers = "*"
+streamlit = "*"
+pandas = "*"
+sentence-transformers = "*"
 
 [tool.poetry.dev-dependencies]
-black = "^22.3.0"
-isort = "^5.10.1"
-mypy = "^0.961"
-flake8 = "^4.0.1"
-flake8-bugbear = "^22.4.25"
-pre-commit = "^2.19.0"
+black = "*"
+isort = "*"
+mypy = "*"
+flake8 = "*"
+flake8-bugbear = "*"
+pre-commit = "*"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]