fix: cache, no image error

app.py

@@ -9,13 +9,20 @@ from transformers import VisionTextDualEncoderModel, VisionTextDualEncoderProces
 st.title("VitB32 Bert Ko Small Clip Test")
 st.markdown("Unsplash data에서 입력 텍스트와 가장 유사한 이미지를 검색합니다.")
 
-with st.spinner("Loading model..."):
-    model = VisionTextDualEncoderModel.from_pretrained(
-        "Bingsu/vitB32_bert_ko_small_clip"
-    )
-    processor = VisionTextDualEncoderProcessor.from_pretrained(
-        "Bingsu/vitB32_bert_ko_small_clip"
-    )
+
+@st.cache(allow_output_mutation=True)
+def get_model():
+    with st.spinner("Loading model..."):
+        model = VisionTextDualEncoderModel.from_pretrained(
+            "Bingsu/vitB32_bert_ko_small_clip"
+        )
+        processor = VisionTextDualEncoderProcessor.from_pretrained(
+            "Bingsu/vitB32_bert_ko_small_clip"
+        )
+    return model, processor
+
+
+model, processor = get_model()
 
 info = pd.read_csv("info.csv")
 with open("img_id.pkl", "rb") as f:
@@ -28,13 +35,23 @@ tokens = processor(text=text, return_tensors="pt")
 with st.spinner("Predicting..."):
     text_emb = model.get_text_features(**tokens)
 
-result = semantic_search(text_emb, img_emb, top_k=6)[0]
+result = semantic_search(text_emb, img_emb, top_k=15)[0]
+_result = iter(result)
+
+
+def get_url() -> str:
+    # 몇몇 이미지가 info.csv 데이터에 없습니다.
+    while True:
+        r = next(_result)
+        photo_id = img_id[r["corpus_id"]]
+        target_series = info.loc[info["photo_id"] == photo_id, "photo_image_url"]
+        if len(target_series) == 0:
+            continue
+        img_url = target_series.iloc[0]
+        return img_url
+
 
 columns = st.columns(3) + st.columns(3)
-for i, col in enumerate(columns):
-    photo_id = img_id[result[i]["corpus_id"]]
-    target_series = info.loc[info["photo_id"] == photo_id, "photo_image_url"]
-    if len(target_series) == 0:
-        continue
-    img_url = target_series.iloc[0]
+for col in columns:
+    img_url = get_url()
     col.image(img_url, use_column_width=True)