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. For 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 describing a function we want to create and search what function in Github already is similar to our search. Right here!:

**Disclaimer**: We will use a sample dataset of 200,000 Python functions on Github for agility. 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 do 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 function description from the code search net dataset (the 'func_documentation_string' column) into a vector space. That is, convert function descriptions from words to numbers to understand 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 description (in this case: 'Create a dictionary') **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): our original text database and our query (here, the description of a python function). Our goal: get the texts in the database that have the closest meaning to our query.

Most similar queries 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 our search's most similar Github function.

"""
)

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] [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 these 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))
"""
)