Add radio button to select target vectors

app.py

@@ -298,10 +298,6 @@ class ClipModel(nn.Module):
         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)
     text_embedding = model.text_model.encode_text([text]).numpy()
@@ -320,8 +316,6 @@ description_text = st.empty()
 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 = DummyClipModel(text_model)
     model = ClipModel("sonoisa/clip-vit-b-32-japanese-v1", device=device)
     st.session_state.model = model
 
@@ -358,6 +352,8 @@ query_input = st.text_input(label="説明文", value="", on_change=clear_result)
 
 closest_n = st.number_input(label="検索数", min_value=1, value=10, max_value=100) 
 
+model_type = st.radio(label="検索対象ベクトル", options=("文", "画像"))
+
 search_buttion = st.button("検索")
 
 result_text = st.empty()
@@ -366,6 +362,11 @@ if search_buttion or prev_query != query_input:
     prev_query = query_input
     query_embedding = encode_text(query_input, model)
 
+    if model_type == "画像":
+        target_vectors = image_vectors
+    else:
+        target_vectors = sentence_vectors
+
     distances = scipy.spatial.distance.cdist(
         query_embedding, image_vectors, metric="cosine"
     )[0]