feat: include new model

app.py

@@ -14,32 +14,35 @@ st.markdown(
 
 
 @st.cache(allow_output_mutation=True, show_spinner=False)
-def get_dual_encoder_model():
+def get_dual_encoder_model(name: str):
     with st.spinner("Loading model..."):
-        model = AutoModel.from_pretrained("Bingsu/vitB32_bert_ko_small_clip").eval()
-        processor = AutoProcessor.from_pretrained("Bingsu/vitB32_bert_ko_small_clip")
+        model = AutoModel.from_pretrained(name).eval()
+        processor = AutoProcessor.from_pretrained(name)
     return model, processor
 
 
 @st.cache(allow_output_mutation=True, show_spinner=False)
-def get_clip_model():
+def get_clip_model(name: str):
     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")
+        model = AutoModel.from_pretrained(name).eval()
+        processor = AutoProcessor.from_pretrained(name)
     return model, processor
 
 
-model_type = st.radio(
-    "Select model",
-    ["Bingsu/clip-vit-base-patch32-ko", "Bingsu/vitB32_bert_ko_small_clip"],
-)
+model_list = [
+    "Bingsu/clip-vit-base-patch32-ko",
+    "Bingsu/clip-vit-large-patch14-ko",
+    "openai/clip-vit-base-patch32",
+    "openai/clip-vit-base-patch16",
+    "openai/clip-vit-large-patch14",
+    "Bingsu/vitB32_bert_ko_small_clip",
+]
+model_type = st.radio("Select model", model_list)
 
-if model_type == "Bingsu/clip-vit-base-patch32-ko":
-    model, processor = get_clip_model()
-elif model_type == "Bingsu/vitB32_bert_ko_small_clip":
-    model, processor = get_dual_encoder_model()
+if model_type == "Bingsu/vitB32_bert_ko_small_clip":
+    model, processor = get_dual_encoder_model(model_type)
 else:
-    raise ValueError("Invalid model type")
+    model, processor = get_clip_model(model_type)
 
 info = pd.read_csv("info.csv")
 with open("img_id.pkl", "rb") as f:
@@ -52,7 +55,7 @@ tokens = processor(text=text, return_tensors="pt")
 with torch.no_grad():
     text_emb = model.get_text_features(**tokens)
 
-result = semantic_search(text_emb, img_emb, top_k=15)[0]
+result = semantic_search(text_emb, img_emb, top_k=16)[0]
 _result = iter(result)