<jupyter_start><jupyter_text>Multilingual Semantic Search with Faiss  Outline- [0. Setup](0)- [1. Load dataset](1)- [2. Build the semantic-similarity search engine](2)- [3. Test the semantic-similarity search engine](3)- [4. References](4)  0. Setup<jupyter_code>#@title Setup Environment
#@markdown Install [Faiss](https://github.com/facebookresearch/faiss) a library 
#@markdown for efficient similarity search and clustering of dense vectors. 
%%capture
!pip install tensorflow_text ## needed????
!pip install bokeh
!pip install faiss-gpu
!pip install tqdm
#@title Imports & Visualization utils
%load_ext google.colab.data_table
import os

import bokeh
import bokeh.models
import bokeh.plotting
import faiss
import numpy as np
import pandas as pd
import sklearn.metrics.pairwise
import tensorflow.compat.v2 as tf
import tensorflow_hub as hub
from google.colab import data_table
from tensorflow_text import SentencepieceTokenizer
from tqdm import tqdm, trange

def visualize_similarity(
    embeddings_1,
    embeddings_2,
    labels_1,
    labels_2,
    plot_title,
    plot_width=1200,
    plot_height=600,
    xaxis_font_size="12pt",
    yaxis_font_size="12pt",
):

    assert len(embeddings_1) == len(labels_1)
    assert len(embeddings_2) == len(labels_2)

    # arccos based text similarity (Yang et al. 2019; Cer et al. 2019)
    sim = (
        1
        - np.arccos(
            sklearn.metrics.pairwise.cosine_similarity(
                embeddings_1, embeddings_2
            )
        )
        / np.pi
    )

    embeddings_1_col, embeddings_2_col, sim_col = [], [], []
    for i in range(len(embeddings_1)):
        for j in range(len(embeddings_2)):
            embeddings_1_col.append(labels_1[i])
            embeddings_2_col.append(labels_2[j])
            sim_col.append(sim[i][j])
    df = pd.DataFrame(
        zip(embeddings_1_col, embeddings_2_col, sim_col),
        columns=["embeddings_1", "embeddings_2", "sim"],
    )

    mapper = bokeh.models.LinearColorMapper(
        palette=[*reversed(bokeh.palettes.YlOrRd[9])],
        low=df.sim.min(),
        high=df.sim.max(),
    )

    p = bokeh.plotting.figure(
        title=plot_title,
        x_range=labels_1,
        x_axis_location="above",
        y_range=[*reversed(labels_2)],
        plot_width=plot_width,
        plot_height=plot_height,
        tools="save",
        toolbar_location="below",
        tooltips=[("pair", "@embeddings_1 ||| @embeddings_2"), ("sim", "@sim")],
    )
    p.rect(
        x="embeddings_1",
        y="embeddings_2",
        width=1,
        height=1,
        source=df,
        fill_color={"field": "sim", "transform": mapper},
        line_color=None,
    )

    p.title.text_font_size = "12pt"
    p.axis.axis_line_color = None
    p.axis.major_tick_line_color = None
    p.axis.major_label_standoff = 16
    p.xaxis.major_label_text_font_size = xaxis_font_size
    p.xaxis.major_label_orientation = 0.25 * np.pi
    p.yaxis.major_label_text_font_size = yaxis_font_size
    p.min_border_right = 300

    bokeh.io.output_notebook()
    bokeh.io.show(p)

!pip list | grep "bokeh\|faiss\|tensorflow-hub\|tensorflow-text"
module_url = "https://tfhub.dev/google/universal-sentence-encoder-multilingual/3"
model = hub.load(module_url)

def embed_text(input):
    return model(input)<jupyter_output><empty_output><jupyter_text>1. Load dataset We will build a semantic-search index of about 200,000 sentences from a wikipedia corpus. About half be in English and the other half in Spanish to demostrate the multilingual capabilities of the Universal Sentence Encoder.First, we will download news sentences in multiple languages from the News Comentary Corpus.<jupyter_code>corpus_metadata = [
    ("es", "Spanish", "en-es.txt.zip", "News-Commentary.en-es.es"),
    ("en", "English", "en-es.txt.zip", "News-Commentary.en-es.en"),
]

sentences = []
language_to_sentences = {}
language_to_news_path = {}

corpus_url = "http://opus.nlpl.eu/download.php?f=News-Commentary/v11/moses/"

for language_code, language_name, zip_file, news_file in corpus_metadata:
    zip_path = tf.keras.utils.get_file(
        fname=zip_file, origin=f"{corpus_url}{zip_file}", extract=True
    )
    news_path = os.path.join(os.path.dirname(zip_path), news_file)
    language_to_sentences[language_code] = pd.read_csv(
        news_path, sep="\t", header=None
    )[0]
    language_to_news_path[language_code] = news_path
    sentences.extend(language_to_sentences[language_code])
    print(f"{len(language_to_sentences[language_code]):,} {language_name}")
explore_df = pd.DataFrame(
    {
        "en": language_to_sentences["en"][:1000],
        "es": language_to_sentences["es"][:1000],
    }
)
data_table.DataTable(explore_df, include_index=False, num_rows_per_page=5)<jupyter_output><empty_output><jupyter_text>2. Build the semantic-similarity search engine  Using a pre-trained model to transform sentences into vectorsWe compute embeddings in _batches_ so that they fit in the GPU's RAM.<jupyter_code>batch_size = 2048
embeddings = []
for indx in trange(0, len(sentences), batch_size):
    embeddings.extend(embed_text(sentences[indx:indx+batch_size]))
    
embeddings = np.array(embeddings).astype("float32")<jupyter_output>100%|██████████| 234/234 [01:27<00:00,  2.69it/s]<jupyter_text>Building an index of semantic vectors<jupyter_code>num_sentences, embedding_dimensions = embeddings.shape
print(f"Number of sentences:  {num_sentences:,}")
print(f"Embedding dimensions: {embedding_dimensions}")
index = faiss.IndexFlatL2(embedding_dimensions)
index = faiss.IndexIDMap(index)
ids = np.array(list(range(num_sentences)))
index.add_with_ids(embeddings, ids)
print(f"Number of vectors in the Faiss index: {index.ntotal}")
D, I = index.search(np.array([embeddings[2]]), k=10)
print(
    f"L2 distance: {D.flatten().tolist()}\n\nMAG sentence ids: {I.flatten().tolist()}"
)
pd.DataFrame({"Similar Sentences": np.array(sentences)[I].flatten()})<jupyter_output><empty_output><jupyter_text>3. Test the search engine- Try different enter sentences or select from the samples list- Try different languages (EN/ES)- Try changing the number of similar sentences<jupyter_code>sample_query = "Global warming" #@param ["Global warming", "Researchers made a surprising new discovery last week.", "The stock market fell four points.", "Lawmakers will vote on the proposal tomorrow."] {allow-input: true}
num_results = 10 #@param {type:"slider", min:0, max:100, step:10}

query_embedding = embed_text(sample_query)[0]
_, I = index.search(np.array([query_embedding]), k=num_results)

pd.DataFrame({"Similar Sentences": np.array(sentences)[I].flatten()})<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12494368523802372, 'fraction_numerical': 0.02042348701006157, 'mean_word_length': 4.298329355608592, 'pattern_counts': {'":': 6, '<': 21, '<?xml version=': 0, '>': 20, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '123306', 'license': 'Apache-2.0', 'n_tokens_mistral': 2392, 'n_tokens_neox': 2257, 'n_words': 475, 'path': 'nlp/multilingual/multilingual-semantic-search-faiss.ipynb', 'repo_name': 'martin-fabbri/colab-notebooks', 'size': '6659'}