import streamlit as st
from PIL import Image

from backend import inference

st.title("Using Sentence Transformers for semantic search")

logos_png = Image.open("images/logos.png")
st.image(
    logos_png,
    caption="The cool kids on the block.",
)


st.subheader("Code Search: An Introduction to Semantic Search")

st.markdown(
    """
Suppose you have a database of texts and you want to find which entries best match the meaning of a single text you have. Example, we want to see which function on Github best matches a description of a function we want to create in Python. However, just word-for-word matching will not work due to the complexity of the programming questions. We need to do a "smart" search: that's what semantic search is for.

*"**Semantic search seeks to improve search accuracy by understanding the content of the search query. In contrast to traditional search engines which only find documents based on lexical matches, semantic search can also find synonyms.**" - [SBert Documentation](https://www.sbert.net/examples/applications/semantic-search/README.html)*.


Let's make this interactive and use the power of **[Streamlit](https://docs.streamlit.io/) + [HF Spaces](https://huggingface.co/spaces) + the '[sentence-transformers/all-mpnet-base-v2](https://huggingface.co/sentence-transformers/all-mpnet-base-v2)' model + the [Code Search Net Dataset](https://huggingface.co/datasets/code_search_net) in Hugging Face Datasets** by giving the description of a function we want to create and search what function in Github already is similar to our search. Right here!:

Disclaimer: For agility, we will use a sample dataset of just 10,000 Python functions on Github. These are very few observations for such a broad search. Therefore, the results may not be optimal.
"""
)

anchor = st.text_input(
    "Please enter here the description of a Python function you want to create, we will look for similar ones in Github.",
    value="Create a dictionary",
)

n_texts = st.number_input(
    f"""How many similar functions you want?""", value=2, min_value=1
)
if st.button("Find them."):

    st.markdown(
        """
        **Amazing! These are the functions in the code search net dataset that better match the language of our specific description.** The 'func_documentation_string' column shows the documentation description of the function in Github, the 'repository_name' column shows the name of the repository where the function is, and 'func_code_url' takes you the function.
        """
    )

    st.table(
        inference.query_search(
            anchor, n_texts, "sentence-transformers/all-mpnet-base-v2"
        )
    )


st.subheader("What happens underneath? Obtaining embeddings")

st.markdown(
    f"""
 First we embed our text database: we convert each of the function descriptions from the code search net dataset (the 'func_documentation_string' column) into a vector space. That is, convert function descriptions from words to numbers that provide an understanding of the language within each description. The understanding of each text will be reflected in a vector called embedding.

 The system would look like the following figure. We embed (1) our texts database (in this case, the Github set of function descriptions), and (2) our own description (in this case: '{anchor}') **with the same model**. Notice that our descriptions could be a matrix of several elements.

"""
)


sentenceEmbeddingsProcess_png = Image.open("images/sentenceEmbeddingsProcess.png")
st.image(
    sentenceEmbeddingsProcess_png,
    caption="Embed the database and your query with the same Sentence Transformers model.",
)


st.subheader("Obtaining the closest observations in the vector space")

st.markdown(
    f"""
We now have two numerical representations of texts (embeddings): one for our original text database and one for our own query (here, the description of a python function). Our goal: get the texts in the database that have the closest meaning to our query.

Queries that are most similar to each other will be closer together in the vector space, and queries that differ most will be farther apart.

The following figure (obtained from the [Sentence Transformers documentation](https://www.sbert.net/examples/applications/semantic-search/README.html)) shows in blue how we would represent the code search net dataset, and in orange your '**{anchor}**' query which is outside the original dataset. The blue dot with the annotation 'Relevant Document' would be the most similar Github function to our search.

"""
)

semanticSearchSBert_png = Image.open("images/semanticSearchSBert.png")
st.image(
    semanticSearchSBert_png,
    caption="Representation of embeddings in a 2D vector space. Obtained from the Sentence Transformers documentation.",
)

st.markdown(
    """

We compare the embedding of our query with the embeddings of each of the texts in the database (there are easier ways to do it but in this case it won't be necessary) using the cosine similarity function, better explained in the [Pytorch documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.cosine_similarity.html). The results of the cosine similarity function will detect which of the texts in the database are closest to our query in vector space.

This is what we did in the dynamic example above!
"""
)


st.subheader("Check this cool resources!")


st.markdown(
    """
If you want to know more about the Sentence Transformers library:
- **The [Hugging Face Organization](https://huggingface.co/sentence-transformers)** for all the new models and instructions on how to download models.
- The **[Nils Reimers tweet](https://twitter.com/Nils_Reimers/status/1487014195568775173) comparing Sentence Transformer models with GPT-3 Embeddings**. Spoiler alert: the Sentence Transformers are awesome!
- The **[Sentence Transformers documentation](https://www.sbert.net/)**,
"""
)

separator = """
---
"""

st.markdown(separator)

st.markdown(
    """
Thanks for reading!

- Omar Espejel ([@espejelomar](https://twitter.com/espejelomar))
"""
)