Sync with data tooling repo, using edugp/kenlm models, updating viz to use quantiles for coloring and ad-hoc viz for the registry dataset

README.md

@@ -11,7 +11,11 @@ pinned: false
 # Installation:
 Requires Python >= 3.7 and < 3.10
 ```
-pip install -r requirements.txt
+pip install .
+```
+Or with [poetry](https://python-poetry.org/)
+```
+poetry install
 ```
 
 # Web App:

app.py

@@ -3,16 +3,28 @@ from functools import partial
 
 import streamlit as st
 from embedding_lenses.data import uploaded_file_to_dataframe
-from embedding_lenses.dimensionality_reduction import (get_tsne_embeddings,
-                                                       get_umap_embeddings)
+from embedding_lenses.dimensionality_reduction import (
+    get_tsne_embeddings,
+    get_umap_embeddings,
+)
 from embedding_lenses.embedding import load_model
 
-from perplexity_lenses.data import (documents_df_to_sentences_df,
-                                    hub_dataset_to_dataframe)
-from perplexity_lenses.engine import (DIMENSIONALITY_REDUCTION_ALGORITHMS,
-                                      DOCUMENT_TYPES, EMBEDDING_MODELS,
-                                      LANGUAGES, SEED, generate_plot)
+from perplexity_lenses import REGISTRY_DATASET
+from perplexity_lenses.data import (
+    documents_df_to_sentences_df,
+    hub_dataset_to_dataframe,
+)
+from perplexity_lenses.engine import (
+    DIMENSIONALITY_REDUCTION_ALGORITHMS,
+    DOCUMENT_TYPES,
+    EMBEDDING_MODELS,
+    LANGUAGES,
+    PERPLEXITY_MODELS,
+    SEED,
+    generate_plot,
+)
 from perplexity_lenses.perplexity import KenlmModel
+from perplexity_lenses.visualization import draw_histogram
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -43,12 +55,36 @@ with col6:
     doc_type = st.selectbox("Document type", DOCUMENT_TYPES, 1)
 with col7:
     sample = st.number_input("Maximum number of documents to use", 1, 100000, 1000)
+perplexity_model = st.selectbox(
+    "Dataset on which the perplexity model was trained on", PERPLEXITY_MODELS, 0
+).lower()
 
 dimensionality_reduction = st.selectbox(
     "Dimensionality Reduction algorithm", DIMENSIONALITY_REDUCTION_ALGORITHMS, 0
 )
 model_name = st.selectbox("Sentence embedding model", EMBEDDING_MODELS, 0)
 
+advanced_options = st.checkbox(
+    "Advanced options (do not modify if using default KenLM models).", value=False
+)
+lower_case = True
+remove_accents = True
+normalize_numbers = True
+punctuation = 1
+if advanced_options:
+    lower_case = st.checkbox(
+        "Lower case text for KenLM preprocessing (from cc_net)", value=False
+    )
+    remove_accents = st.checkbox(
+        "Remove accents for KenLM preprocessing (from cc_net)", value=False
+    )
+    normalize_numbers = st.checkbox(
+        "Replace numbers with zeros KenLM preprocessing (from cc_net)", value=True
+    )
+    punctuation = st.number_input(
+        "Punctuation mode to use from cc_net KenLM preprocessing", 1, 2, 1
+    )
+
 with st.spinner(text="Loading embedding model..."):
     model = load_model(model_name)
 dimensionality_reduction_function = (
@@ -58,7 +94,14 @@ dimensionality_reduction_function = (
 )
 
 with st.spinner(text="Loading KenLM model..."):
-    kenlm_model = KenlmModel.from_pretrained(language)
+    kenlm_model = KenlmModel.from_pretrained(
+        perplexity_model,
+        language,
+        lower_case,
+        remove_accents,
+        normalize_numbers,
+        punctuation,
+    )
 
 if uploaded_file or hub_dataset:
     with st.spinner("Loading dataset..."):
@@ -84,7 +127,7 @@ if uploaded_file or hub_dataset:
     logger.info(
         f"Perplexity range: {df['perplexity'].min()} - {df['perplexity'].max()}"
     )
-    plot = generate_plot(
+    plot, plot_registry = generate_plot(
         df,
         text_column,
         "perplexity",
@@ -93,7 +136,12 @@ if uploaded_file or hub_dataset:
         model,
         seed=SEED,
         context_logger=st.spinner,
+        hub_dataset=hub_dataset,
     )
-    logger.info("Displaying plot")
+    logger.info("Displaying plots")
     st.bokeh_chart(plot)
+    if hub_dataset == REGISTRY_DATASET:
+        st.bokeh_chart(plot_registry)
+    fig = draw_histogram(df["perplexity"].values)
+    st.pyplot(fig)
     logger.info("Done")

cli.py

@@ -6,17 +6,28 @@ import pandas as pd
 import typer
 from bokeh.plotting import output_file as bokeh_output_file
 from bokeh.plotting import save
-from embedding_lenses.data import uploaded_file_to_dataframe
-from embedding_lenses.dimensionality_reduction import (get_tsne_embeddings,
-                                                       get_umap_embeddings)
+from embedding_lenses.dimensionality_reduction import (
+    get_tsne_embeddings,
+    get_umap_embeddings,
+)
 from embedding_lenses.embedding import load_model
 
-from perplexity_lenses.data import (documents_df_to_sentences_df,
-                                    hub_dataset_to_dataframe)
-from perplexity_lenses.engine import (DIMENSIONALITY_REDUCTION_ALGORITHMS,
-                                      DOCUMENT_TYPES, EMBEDDING_MODELS,
-                                      LANGUAGES, SEED, generate_plot)
+from perplexity_lenses import REGISTRY_DATASET
+from perplexity_lenses.data import (
+    documents_df_to_sentences_df,
+    hub_dataset_to_dataframe,
+)
+from perplexity_lenses.engine import (
+    DIMENSIONALITY_REDUCTION_ALGORITHMS,
+    DOCUMENT_TYPES,
+    EMBEDDING_MODELS,
+    LANGUAGES,
+    PERPLEXITY_MODELS,
+    SEED,
+    generate_plot,
+)
 from perplexity_lenses.perplexity import KenlmModel
+from perplexity_lenses.visualization import draw_histogram
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -46,6 +57,10 @@ def main(
         help=f"Whether to embed at the sentence or document level. Options: {DOCUMENT_TYPES}.",
     ),
     sample: int = typer.Option(1000, help="Maximum number of examples to use."),
+    perplexity_model: str = typer.Option(
+        "wikipedia",
+        help=f"Dataset on which the perplexity model was trained on. Options: {PERPLEXITY_MODELS}",
+    ),
     dimensionality_reduction: str = typer.Option(
         DIMENSIONALITY_REDUCTION_ALGORITHMS[0],
         help=f"Whether to use UMAP or t-SNE for dimensionality reduction. Options: {DIMENSIONALITY_REDUCTION_ALGORITHMS}.",
@@ -55,7 +70,7 @@ def main(
         help=f"The sentence embedding model to use. Options: {EMBEDDING_MODELS}",
     ),
     output_file: str = typer.Option(
-        "perplexity.html", help="The name of the output visualization HTML file."
+        "perplexity", help="The name of the output visualization files."
     ),
 ):
     """
@@ -69,7 +84,14 @@ def main(
         else partial(get_tsne_embeddings, random_state=SEED)
     )
     logger.info("Loading KenLM model...")
-    kenlm_model = KenlmModel.from_pretrained(language)
+    kenlm_model = KenlmModel.from_pretrained(
+        perplexity_model.lower(),
+        language,
+        lower_case=True,
+        remove_accents=True,
+        normalize_numbers=True,
+        punctuation=1,
+    )
     logger.info("Loading dataset...")
     if dataset.endswith(".csv") or dataset.endswith(".tsv"):
         df = pd.read_csv(dataset, sep="\t" if dataset.endswith(".tsv") else ",")
@@ -92,7 +114,7 @@ def main(
     logger.info(
         f"Perplexity range: {df['perplexity'].min()} - {df['perplexity'].max()}"
     )
-    plot = generate_plot(
+    plot, plot_registry = generate_plot(
         df,
         text_column,
         "perplexity",
@@ -100,10 +122,16 @@ def main(
         dimensionality_reduction_function,
         model,
         seed=SEED,
+        hub_dataset=dataset,
     )
-    logger.info("Saving plot")
-    bokeh_output_file(output_file)
+    logger.info("Saving plots")
+    bokeh_output_file(f"{output_file}.html")
     save(plot)
+    if dataset == REGISTRY_DATASET:
+        bokeh_output_file(f"{output_file}_registry.html")
+        save(plot_registry)
+    fig = draw_histogram(df["perplexity"].values)
+    fig.savefig(f"{output_file}_histogram.png")
     logger.info("Done")
 
 

perplexity_lenses/__init__.py

@@ -1 +1,2 @@
 __version__ = "0.1.0"
+REGISTRY_DATASET = "mhtoin/register_oscar"

perplexity_lenses/data.py

@@ -5,6 +5,7 @@ import pandas as pd
 from datasets import load_dataset
 from tqdm import tqdm
 
+from perplexity_lenses import REGISTRY_DATASET
 from perplexity_lenses.perplexity import KenlmModel
 
 
@@ -21,13 +22,22 @@ def hub_dataset_to_dataframe(
     load_dataset_fn = partial(load_dataset, path=path)
     if name:
         load_dataset_fn = partial(load_dataset_fn, name=name)
+        # Special case for the registry dataset
+        if path == REGISTRY_DATASET:
+            load_dataset_fn = partial(load_dataset_fn, data_files=f"{name}/*")
     if split:
         load_dataset_fn = partial(load_dataset_fn, split=split)
     dataset = load_dataset_fn(streaming=True).shuffle(buffer_size=10000, seed=seed)
     if doc_type.lower() == "sentence":
         dataset = dataset.map(
             lambda x: [
-                {text_column: sentence, "perplexity": model.get_perplexity(sentence)}
+                {
+                    text_column: sentence,
+                    "perplexity": model.get_perplexity(sentence),
+                    "label": x.get("labels", [])[0]
+                    if len(x.get("labels", [])) > 0
+                    else "NONE",  # Special case for registry dataset
+                }
                 for sentence in x[text_column].split("\n")
             ]
         )
@@ -36,6 +46,9 @@ def hub_dataset_to_dataframe(
             lambda x: {
                 text_column: x[text_column],
                 "perplexity": model.get_perplexity(x[text_column]),
+                "label": x.get("labels", [])[0]
+                if len(x.get("labels", [])) > 0
+                else "NONE",  # Special case for registry dataset
             }
         )
     instances = []

perplexity_lenses/engine.py

@@ -1,15 +1,17 @@
 import logging
 import time
-from typing import Callable, Optional, Union
+from typing import Callable, Optional, Tuple, Union
 
 import pandas as pd
 import streamlit as st
+from bokeh.palettes import Turbo256
 from bokeh.plotting import Figure
 from embedding_lenses.embedding import embed_text
 from embedding_lenses.utils import encode_labels
+from embedding_lenses.visualization import draw_interactive_scatter_plot
 from sentence_transformers import SentenceTransformer
 
-from perplexity_lenses.visualization import draw_interactive_scatter_plot
+from perplexity_lenses import REGISTRY_DATASET
 
 logger = logging.getLogger(__name__)
 EMBEDDING_MODELS = [
@@ -71,6 +73,7 @@ LANGUAGES = [
     "uk",
     "zh",
 ]
+PERPLEXITY_MODELS = ["Wikipedia", "OSCAR"]
 
 
 class ContextLogger:
@@ -94,7 +97,8 @@ def generate_plot(
     model: SentenceTransformer,
     seed: int = 0,
     context_logger: Union[st.spinner, ContextLogger] = ContextLogger,
-) -> Figure:
+    hub_dataset: str = "",
+) -> Tuple[Figure, Optional[Figure]]:
     if text_column not in df.columns:
         raise ValueError(
             f"The specified column name doesn't exist. Columns available: {df.columns.values}"
@@ -111,13 +115,32 @@ def generate_plot(
     with context_logger("Reducing dimensionality..."):
         embeddings_2d = dimensionality_reduction_function(embeddings)
     logger.info("Generating figure")
+    hover_data = {
+        text_column: df[text_column].values,
+        label_column: encoded_labels.values,
+    }
+    # Round perplexity values
+    values = df[label_column].values.round().astype(int)
     plot = draw_interactive_scatter_plot(
-        df[text_column].values,
+        hover_data,
         embeddings_2d[:, 0],
         embeddings_2d[:, 1],
-        encoded_labels.values,
-        df[label_column].values,
-        text_column,
-        label_column,
+        values,
     )
-    return plot
+    # Special case for the registry dataset
+    plot_registry = None
+    if hub_dataset == REGISTRY_DATASET:
+        encoded_labels = encode_labels(df["label"])
+        hover_data = {
+            text_column: df[text_column].values,
+            "label": df["label"].values,
+            label_column: df[label_column].values,
+        }
+        plot_registry = draw_interactive_scatter_plot(
+            hover_data,
+            embeddings_2d[:, 0],
+            embeddings_2d[:, 1],
+            encoded_labels.values,
+            palette=Turbo256,
+        )
+    return plot, plot_registry

perplexity_lenses/perplexity.py

@@ -1,11 +1,13 @@
 import os
 import re
 import unicodedata
-import urllib.request
 from typing import Dict
 
 import kenlm
 import sentencepiece
+from huggingface_hub import cached_download, hf_hub_url
+
+KENLM_MODEL_REPO = "edugp/kenlm"
 
 
 class SentencePiece:
@@ -64,30 +66,65 @@ class KenlmModel:
     non_printing_chars_re = re.compile(
         f"[{''.join(map(chr, list(range(0,32)) + list(range(127,160))))}]"
     )
+    kenlm_model_dir = None
+    sentence_piece_model_dir = None
 
-    def __init__(self, language):
-        download_kenlm_model(language)
+    def __init__(
+        self,
+        model_dataset: str,
+        language: str,
+        lower_case: bool = False,
+        remove_accents: bool = False,
+        normalize_numbers: bool = True,
+        punctuation: int = 1,
+    ):
+        self.download_kenlm_model(model_dataset, language)
         try:
-            self.model = kenlm.Model(f"{language}.arpa.bin")
-            self.tokenizer = SentencePiece(f"{language}.sp.model")
+            self.model = kenlm.Model(self.kenlm_model_dir)
+            self.tokenizer = SentencePiece(self.sentence_piece_model_dir)
         except OSError:
-            os.remove(f"{language}.arpa.bin")
-            if os.path.exists(f"{language}.sp.model"):
-                os.remove(f"{language}.sp.model")
+            os.remove(self.kenlm_model_dir)
+            if os.path.exists(self.sentence_piece_model_dir):
+                os.remove(self.sentence_piece_model_dir)
             raise OSError(
                 "File was corrupt and should have been removed. Please, retry."
             )
+        self.accent = remove_accents
+        self.case = lower_case
+        self.numbers = normalize_numbers
+        self.punct = punctuation
 
     @classmethod
-    def from_pretrained(cls, language: str):
-        return cls(language)
+    def from_pretrained(
+        cls,
+        model_dataset: str,
+        language: str,
+        lower_case: bool,
+        remove_accents: bool,
+        normalize_numbers: bool,
+        punctuation: int,
+    ):
+        return cls(
+            model_dataset,
+            language,
+            lower_case,
+            remove_accents,
+            normalize_numbers,
+            punctuation,
+        )
 
     def pp(self, log_score, length):
         return 10.0 ** (-log_score / length)
 
     def get_perplexity(self, doc: str, normalize_cc_net: bool = True):
         if normalize_cc_net:
-            doc = self.normalize(doc)
+            doc = self.normalize(
+                doc,
+                accent=self.accent,
+                case=self.case,
+                numbers=self.numbers,
+                punct=self.punct,
+            )
         # Tokenize (after normalizing): See https://github.com/facebookresearch/cc_net/blob/bda555bd1cf1ee2e0b925363e62a61cd46c8b60d/cc_net/mine.py#L352 for full pipeline
         doc = self.tokenizer.do(doc)
         doc_log_score, doc_length = 0, 0
@@ -140,16 +177,12 @@ class KenlmModel:
     def remove_non_printing_char(self, text: str) -> str:
         return self.non_printing_chars_re.sub("", text)
 
-
-def download_kenlm_model(language: str):
-    root_url = "http://dl.fbaipublicfiles.com/cc_net/lm"
-    bin_name = f"{language}.arpa.bin"
-    model_name = f"{language}.sp.model"
-    bin_url = f"{root_url}/{bin_name}"
-    model_url = f"{root_url}/{model_name}"
-
-    if not os.path.isfile(bin_name):
-        urllib.request.urlretrieve(bin_url, bin_name)
-
-    if not os.path.isfile(model_name):
-        urllib.request.urlretrieve(model_url, model_name)
+    def download_kenlm_model(self, model_dataset: str, language: str):
+        kenlm_model_url = hf_hub_url(
+            KENLM_MODEL_REPO, filename=f"{model_dataset}/{language}.arpa.bin"
+        )
+        self.kenlm_model_dir = cached_download(kenlm_model_url)
+        sentence_piece_model_url = hf_hub_url(
+            KENLM_MODEL_REPO, filename=f"{model_dataset}/{language}.sp.model"
+        )
+        self.sentence_piece_model_dir = cached_download(sentence_piece_model_url)

perplexity_lenses/visualization.py

@@ -1,57 +1,18 @@
+import matplotlib.figure
+import matplotlib.pyplot as plt
 import numpy as np
-from bokeh.models import ColumnDataSource, HoverTool
-from bokeh.palettes import Cividis256 as Pallete
-from bokeh.plotting import Figure, figure
-from bokeh.transform import factor_cmap
 
 
-def draw_interactive_scatter_plot(
-    texts: np.ndarray,
-    xs: np.ndarray,
-    ys: np.ndarray,
+def draw_histogram(
     values: np.ndarray,
-    labels: np.ndarray,
-    text_column: str,
-    label_column: str,
-) -> Figure:
-    # Smooth down values for coloring, by taking the entropy = log10(perplexity) and multiply it by 10000
-    values = ((np.log10(values)) * 10000).round().astype(int)
-    # Normalize values to range between 0-255, to assign a color for each value
-    max_value = values.max()
-    min_value = values.min()
-    if max_value - min_value == 0:
-        values_color = np.ones(len(values))
-    else:
-        values_color = (
-            ((values - min_value) / (max_value - min_value) * 255).round().astype(int)
-        )
-    values_color_sorted = sorted(values_color)
-
-    values_list = values.astype(str).tolist()
-    values_sorted = sorted(values_list)
-    labels_list = labels.astype(str).tolist()
-
-    source = ColumnDataSource(
-        data=dict(x=xs, y=ys, text=texts, label=values_list, original_label=labels_list)
-    )
-    hover = HoverTool(
-        tooltips=[(text_column, "@text{safe}"), (label_column, "@original_label")]
-    )
-    p = figure(plot_width=800, plot_height=800, tools=[hover])
-    p.circle(
-        "x",
-        "y",
-        size=10,
-        source=source,
-        fill_color=factor_cmap(
-            "label",
-            palette=[Pallete[id_] for id_ in values_color_sorted],
-            factors=values_sorted,
-        ),
-    )
-
-    p.axis.visible = False
-    p.xgrid.grid_line_color = None
-    p.ygrid.grid_line_color = None
-    p.toolbar.logo = None
-    return p
+    cutoff_x_axis: float = 2000.0,
+    title: str = "Perplexity histogram",
+    xlabel: str = "Perplexity",
+) -> matplotlib.figure.Figure:
+    hist_values = values[values < cutoff_x_axis]
+    fig, ax = plt.subplots(figsize=(12, 9))
+    ax.hist(hist_values, bins=50)
+    ax.set_title(title)
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel("Counts")
+    return fig

requirements.txt

@@ -1,5 +1,5 @@
 bokeh==2.2.2
-https://files.pythonhosted.org/packages/2f/58/e00d2495b54f4ba97ca31a11aa7e636f80183ccf9b616f7eaa5518d050bb/embedding_lenses-0.5.0-py3-none-any.whl
+https://files.pythonhosted.org/packages/52/b6/798b1ac755c19cc5a91e5f5068544ed9deda74cb9339e4bb5e2b5f3ff3ec/embedding_lenses-0.8.0-py3-none-any.whl
 https://github.com/kpu/kenlm/archive/master.zip
 huggingface-hub==0.0.19
 numpy==1.20.0