import streamlit as st
import numpy as np
from html import escape
import torch
from transformers import AutoModel, AutoTokenizer

text_tokenizer = AutoTokenizer.from_pretrained('arman-aminian/farsi-image-search-text')
text_encoder = AutoModel.from_pretrained('arman-aminian/farsi-image-search-text').eval()
image_embeddings = torch.load('image_embeddings.pt')
image_links = np.load('image_links.npy', allow_pickle=True)


@st.experimental_memo
def image_search(query, top_k=10):
    with torch.no_grad():
        text_embedding = text_encoder(**text_tokenizer(query, return_tensors='pt')).pooler_output
    _, indices = torch.cosine_similarity(image_embeddings, text_embedding).sort(descending=True)
    return [image_links[i] for i in indices[:top_k]]


def get_html(url_list):
    html = "<div style='margin-top: 50px; max-width: 1100px; display: flex; flex-wrap: wrap; justify-content: space-evenly'>"
    for url in url_list:
        html2 = f"<img style='height: 180px; margin: 2px' src='{escape(url)}'>"  
        html = html + html2
    html += "</div>"
    return html


def main():
    st.markdown('''
                <style>
                .block-container{
                  max-width: 1200px;
                }
                section.main>div:first-child {
                  padding-top: 0px;
                }
                section:not(.main)>div:first-child {
                  padding-top: 30px;
                }
                div.reportview-container > section:first-child{
                  max-width: 320px;
                }
                #MainMenu {
                  visibility: hidden;
                }
                footer {
                  visibility: hidden;
                }
                </style>''',
            
                unsafe_allow_html=True)        
         
                
    st.sidebar.markdown('''
# CLIP FARSI
Enter the query. We search for you among [25 thousand photos](https://unsplash.com/) and bring you the most relevant photos.
The thing about this model is that it searches for you among the raw photos and these photos do not have explanations next to them. For a better understanding, we will give some examples of the applications that this model can have
- For example, you can search in Farsi in your phone's photo gallery
- When you are writing a blog or any text, you can put related photos between your texts. In this way, you give the paragraph for which you want a photo to the model, so that it will find a photo related to it for you
Note: We used a small collection (25 thousand) of images to keep this program real-time, but obviously the quality of the image search depends heavily on the size of the image database and this version is just an initial demo to familiarize you with the model.
''')
    _, c, _ = st.columns((1, 3, 1))
    query = c.text_input('Search here')
    if len(query) > 0:
        results = image_search(query)
        st.markdown(get_html(results), unsafe_allow_html=True)


if __name__ == '__main__':
    main()