Support visualizing both sentences and whole documents. Smooth down color assignment in visualization.

app.py

@@ -6,15 +6,14 @@ import pandas as pd
 import streamlit as st
 from bokeh.plotting import Figure
 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 embed_text, load_model
 from embedding_lenses.utils import encode_labels
-from embedding_lenses.visualization import draw_interactive_scatter_plot
 from sentence_transformers import SentenceTransformer
 
-from data import hub_dataset_to_dataframe
-from perplexity import KenlmModel
+from perplexity_lenses.data import documents_df_to_sentences_df, hub_dataset_to_dataframe
+from perplexity_lenses.perplexity import KenlmModel
+from perplexity_lenses.visualization import draw_interactive_scatter_plot
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -70,6 +69,7 @@ LANGUAGES = [
     "uk",
     "zh",
 ]
+DOCUMENT_TYPES = ["Whole document", "Sentence"]
 SEED = 0
 
 
@@ -113,9 +113,18 @@ with col2:
 with col3:
     hub_dataset_split = st.text_input("Dataset split", "train")
 
-text_column = st.text_input("Text field name", "text")
-language = st.selectbox("Language", LANGUAGES, 12)
-sample = st.number_input("Maximum number of documents to use", 1, 100000, 1000)
+col4, col5 = st.columns(2)
+with col4:
+    text_column = st.text_input("Text field name", "text")
+with col5:
+    language = st.selectbox("Language", LANGUAGES, 12)
+
+col6, col7 = st.columns(2)
+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)
+
 dimensionality_reduction = st.selectbox("Dimensionality Reduction algorithm", DIMENSIONALITY_REDUCTION_ALGORITHMS, 0)
 model_name = st.selectbox("Sentence embedding model", EMBEDDING_MODELS, 0)
 
@@ -132,10 +141,16 @@ if uploaded_file or hub_dataset:
     with st.spinner("Loading dataset..."):
         if uploaded_file:
             df = uploaded_file_to_dataframe(uploaded_file)
+            if doc_type == "Sentence":
+                df = documents_df_to_sentences_df(df, text_column, sample, seed=SEED)
             df["perplexity"] = df[text_column].map(kenlm_model.get_perplexity)
         else:
-            df = hub_dataset_to_dataframe(hub_dataset, hub_dataset_config, hub_dataset_split, sample, text_column, kenlm_model, seed=SEED)
-    plot = generate_plot(df, text_column, "perplexity", sample, dimensionality_reduction_function, model)
+            df = hub_dataset_to_dataframe(hub_dataset, hub_dataset_config, hub_dataset_split, sample, text_column, kenlm_model, seed=SEED, doc_type=doc_type)
+
+    # Round perplexity
+    df["perplexity"] = df["perplexity"].round().astype(int)
+    logger.info(f"Perplexity range: {df['perplexity'].min()} - {df['perplexity'].max()}")
+    plot = generate_plot(df, text_column, "perplexity", None, dimensionality_reduction_function, model)
     logger.info("Displaying plot")
     st.bokeh_chart(plot)
     logger.info("Done")

data.py

@@ -1,28 +0,0 @@
-from functools import partial
-
-import pandas as pd
-from datasets import load_dataset
-from tqdm import tqdm
-
-from perplexity import KenlmModel
-
-
-def hub_dataset_to_dataframe(path: str, name: str, split: str, sample: int, text_column: str, model: KenlmModel, seed: int = 0) -> pd.DataFrame:
-    load_dataset_fn = partial(load_dataset, path=path)
-    if name:
-        load_dataset_fn = partial(load_dataset_fn, name=name)
-    if split:
-        load_dataset_fn = partial(load_dataset_fn, split=split)
-    dataset = (
-        load_dataset_fn(streaming=True)
-        .shuffle(buffer_size=10000, seed=seed)
-        .map(lambda x: {text_column: x[text_column], "perplexity": model.get_perplexity(x[text_column])})
-    )
-    instances = []
-    count = 0
-    for instance in tqdm(dataset, total=sample):
-        instances.append(instance)
-        count += 1
-        if count == sample:
-            break
-    return pd.DataFrame(instances)

perplexity_lenses/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

perplexity_lenses/data.py

@@ -0,0 +1,43 @@
+from functools import partial
+
+import numpy as np
+import pandas as pd
+from datasets import load_dataset
+from tqdm import tqdm
+
+from perplexity_lenses.perplexity import KenlmModel
+
+
+def hub_dataset_to_dataframe(
+    path: str, name: str, split: str, sample: int, text_column: str, model: KenlmModel, seed: int = 0, doc_type: str = "Whole document"
+) -> pd.DataFrame:
+    load_dataset_fn = partial(load_dataset, path=path)
+    if name:
+        load_dataset_fn = partial(load_dataset_fn, name=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 == "Sentence":
+        dataset = dataset.map(lambda x: [{text_column: sentence, "perplexity": model.get_perplexity(sentence)} for sentence in x[text_column].split("\n")])
+    else:
+        dataset = dataset.map(lambda x: {text_column: x[text_column], "perplexity": model.get_perplexity(x[text_column])})
+    instances = []
+    count = 0
+    for instance in tqdm(dataset, total=sample):
+        if isinstance(instance, list):
+            for sentence in instance:
+                instances.append(sentence)
+                count += 1
+                if count == sample:
+                    break
+        else:
+            instances.append(instance)
+            count += 1
+        if count == sample:
+            break
+    return pd.DataFrame(instances)
+
+
+def documents_df_to_sentences_df(df: pd.DataFrame, text_column: str, sample: int, seed: int = 0):
+    df_sentences = pd.DataFrame({text_column: np.array(df[text_column].map(lambda x: x.split("\n")).values.tolist()).flatten()})
+    return df_sentences.sample(min(sample, df.shape[0]), random_state=seed)

perplexity_lenses/visualization.py

@@ -0,0 +1,35 @@
+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, 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