fix: norm

app.py

@@ -3,6 +3,7 @@ import pickle
 import numpy as np
 import pandas as pd
 import streamlit as st
+import torch
 from sentence_transformers.util import semantic_search
 from transformers import VisionTextDualEncoderModel, VisionTextDualEncoderProcessor
 
@@ -15,7 +16,7 @@ def get_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"
         )
@@ -31,7 +32,10 @@ img_emb = np.load("img_emb.npy")
 
 text = st.text_input("Input Text", value="검은 고양이")
 tokens = processor(text=text, return_tensors="pt")
-text_emb = model.get_text_features(**tokens)
+
+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)