import pickle

import streamlit as st
import numpy as np
import pandas as pd
from transformers import VisionTextDualEncoderModel, VisionTextDualEncoderProcessor
from sentence_transformers.util import semantic_search


st.title("VitB32 Bert Ko Small Clip Test")

model = VisionTextDualEncoderModel.from_pretrained("Bingsu/vitB32_bert_ko_small_clip")
processor = VisionTextDualEncoderProcessor.from_pretrained(
    "Bingsu/vitB32_bert_ko_small_clip"
)

info = pd.read_csv("info.csv")
with open("img_id.pkl", "rb") as f:
    img_id = pickle.load(f)
img_emb = np.load("img_emb.npy")

text = st.text_input("Input Text")
tokens = processor(text=text, return_tensors="pt")
text_emb = model.get_text_features(**tokens)

result = semantic_search(text_emb, img_emb, top_k=6)

st.write(result)