Add script to generate dataset of embeddings and perplexities. Add script to generate t-SNE plot for embedding and perplexity visualization.

get_embeddings_and_perplexity.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+import kenlm
+from datasets import load_dataset
+from tqdm import tqdm
+import pandas as pd
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+
+TOTAL_SENTENCES = 20000
+def pp(log_score, length):
+    return 10.0 ** (-log_score / length)
+
+
+embedder = "distiluse-base-multilingual-cased-v1"
+embedder_model = SentenceTransformer(embedder)
+embedding_shape = embedder_model.encode(["foo"])[0].shape[0]
+# http://dl.fbaipublicfiles.com/cc_net/lm/es.arpa.bin
+model = kenlm.Model("es.arpa.bin")
+mc4 = load_dataset("mc4", "es", streaming=True)
+count = 0
+embeddings = []
+lenghts = []
+perplexities = []
+sentences = []
+
+for sample in tqdm(mc4["train"].shuffle(buffer_size=100_000), total=416057992):
+    lines = sample["text"].split("\n")
+    for line in lines:
+        count += 1
+        log_score = model.score(line)
+        length = len(line.split()) + 1
+        embedding = embedder_model.encode([line])[0]
+        embeddings.append(embedding.tolist())
+        perplexities.append(pp(log_score, length))
+        lenghts.append(length)
+        sentences.append(line)
+        if count == TOTAL_SENTENCES:
+            break
+    if count == TOTAL_SENTENCES:
+        embeddings = np.array(embeddings)
+        df = pd.DataFrame({"sentence": sentences, "length": lenghts, "perplexity": perplexities})
+        for dim in range(embedding_shape):
+            df[f"dim_{dim}"] = embeddings[:, dim]
+        df.to_csv("mc4-es-perplexity-sentences.tsv", index=None, sep="\t")
+        print("DONE!")
+        break

tsne_plot.py

@@ -0,0 +1,66 @@
+import argparse
+import logging
+from typing import Any, Optional
+
+import bokeh
+import numpy as np
+import pandas as pd
+from bokeh.models import ColumnDataSource, HoverTool
+from bokeh.plotting import figure, output_file, save
+from bokeh.transform import factor_cmap
+from bokeh.palettes import Cividis256 as Pallete
+from bokeh.resources import CDN
+from bokeh.embed import file_html
+from sklearn.manifold import TSNE
+
+
+logging.basicConfig(level = logging.INFO)
+logger = logging.getLogger(__name__)
+SEED = 0
+
+def get_tsne_embeddings(embeddings: np.ndarray, perplexity: int=30, n_components: int=2, init: str='pca', n_iter: int=5000, random_state: int=SEED) -> np.ndarray:
+    tsne = TSNE(perplexity=perplexity, n_components=n_components, init=init, n_iter=n_iter, random_state=random_state)
+    return tsne.fit_transform(embeddings)
+
+def draw_interactive_scatter_plot(texts: np.ndarray, xs: np.ndarray, ys: np.ndarray, values: np.ndarray) -> Any:
+    # Normalize values to range between 0-255, to assign a color for each value
+    max_value = values.max()
+    min_value = values.min()
+    values_color = ((values - min_value) / (max_value - min_value) * 255).round().astype(int).astype(str)
+    values_color_set = sorted(values_color)
+
+    values_list = values.astype(str).tolist()
+    values_set = sorted(values_list)
+
+    source = ColumnDataSource(data=dict(x=xs, y=ys, text=texts, perplexity=values_list))
+    hover = HoverTool(tooltips=[('Sentence', '@text{safe}'), ('Perplexity', '@perplexity')])
+    p = figure(plot_width=1200, plot_height=1200, tools=[hover], title='Sentences')
+    p.circle(
+        'x', 'y', size=10, source=source, fill_color=factor_cmap('perplexity', palette=[Pallete[int(id_)] for id_ in values_color_set], factors=values_set))
+    return p
+
+def generate_plot(tsv: str, output_file_name: str, sample: Optional[int]):
+    logger.info("Loading dataset in memory")
+    df = pd.read_csv(tsv, sep="\t")
+    if sample:
+        df = df.sample(sample, random_state=SEED)
+    logger.info(f"Dataset contains {df.shape[0]} sentences")
+    embeddings = df[sorted([col for col in df.columns if col.startswith("dim")], key=lambda x: int(x.split("_")[-1]))].values
+    logger.info(f"Running t-SNE")
+    tsne_embeddings = get_tsne_embeddings(embeddings)
+    logger.info(f"Generating figure")
+    plot = draw_interactive_scatter_plot(df["sentence"].values, tsne_embeddings[:, 0], tsne_embeddings[:, 1], df["perplexity"].values)
+    output_file(output_file_name)
+    save(plot)
+
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Embeddings t-SNE plot")
+    parser.add_argument("--tsv", type=str, help="Path to tsv file with columns 'text', 'perplexity' and N 'dim_<i> columns for each embdeding dimension.'")
+    parser.add_argument("--output_file", type=str, help="Path to the output HTML file for the interactive plot.", default="perplexity_colored_embeddings.html")
+    parser.add_argument("--sample", type=int, help="Number of sentences to use", default=None)
+
+    args = parser.parse_args()
+    generate_plot(args.tsv, args.output_file, args.sample)