Add CLI and refactor

app.py

@@ -1,104 +1,17 @@
 import logging
 from functools import partial
-from typing import Callable, Optional
 
-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.embedding import embed_text, load_model
-from embedding_lenses.utils import encode_labels
-from sentence_transformers import SentenceTransformer
+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.perplexity import KenlmModel
-from perplexity_lenses.visualization import draw_interactive_scatter_plot
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-EMBEDDING_MODELS = ["distiluse-base-multilingual-cased-v1", "all-mpnet-base-v2", "flax-sentence-embeddings/all_datasets_v3_mpnet-base"]
-DIMENSIONALITY_REDUCTION_ALGORITHMS = ["UMAP", "t-SNE"]
-LANGUAGES = [
-    "af",
-    "ar",
-    "az",
-    "be",
-    "bg",
-    "bn",
-    "ca",
-    "cs",
-    "da",
-    "de",
-    "el",
-    "en",
-    "es",
-    "et",
-    "fa",
-    "fi",
-    "fr",
-    "gu",
-    "he",
-    "hi",
-    "hr",
-    "hu",
-    "hy",
-    "id",
-    "is",
-    "it",
-    "ja",
-    "ka",
-    "kk",
-    "km",
-    "kn",
-    "ko",
-    "lt",
-    "lv",
-    "mk",
-    "ml",
-    "mn",
-    "mr",
-    "my",
-    "ne",
-    "nl",
-    "no",
-    "pl",
-    "pt",
-    "ro",
-    "ru",
-    "uk",
-    "zh",
-]
-DOCUMENT_TYPES = ["Whole document", "Sentence"]
-SEED = 0
-
-
-def generate_plot(
-    df: pd.DataFrame,
-    text_column: str,
-    label_column: str,
-    sample: Optional[int],
-    dimensionality_reduction_function: Callable,
-    model: SentenceTransformer,
-) -> Figure:
-    if text_column not in df.columns:
-        raise ValueError(f"The specified column name doesn't exist. Columns available: {df.columns.values}")
-    if label_column not in df.columns:
-        df[label_column] = 0
-    df = df.dropna(subset=[text_column, label_column])
-    if sample:
-        df = df.sample(min(sample, df.shape[0]), random_state=SEED)
-    with st.spinner(text="Embedding text..."):
-        embeddings = embed_text(df[text_column].values.tolist(), model)
-    logger.info("Encoding labels")
-    encoded_labels = encode_labels(df[label_column])
-    with st.spinner("Reducing dimensionality..."):
-        embeddings_2d = dimensionality_reduction_function(embeddings)
-    logger.info("Generating figure")
-    plot = draw_interactive_scatter_plot(
-        df[text_column].values, embeddings_2d[:, 0], embeddings_2d[:, 1], encoded_labels.values, df[label_column].values, text_column, label_column
-    )
-    return plot
 
 
 st.title("Perplexity Lenses")
@@ -150,7 +63,7 @@ if uploaded_file or hub_dataset:
     # 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)
+    plot = generate_plot(df, text_column, "perplexity", None, dimensionality_reduction_function, model, seed=SEED, context_logger=st.spinner)
     logger.info("Displaying plot")
     st.bokeh_chart(plot)
     logger.info("Done")

cli.py

@@ -0,0 +1,68 @@
+import logging
+from functools import partial
+from typing import Optional
+
+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.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.perplexity import KenlmModel
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+app = typer.Typer()
+
+
+@app.command()
+def main(
+    dataset: str = typer.Option("mc4", help="The name of the hub dataset or local csv/tsv file."),
+    dataset_config: Optional[str] = typer.Option("es", help="The configuration of the hub dataset, if any. Does not apply to local csv/tsv files."),
+    dataset_split: Optional[str] = typer.Option("train", help="The dataset split. Does not apply to local csv/tsv files."),
+    text_column: str = typer.Option("text", help="The text field name."),
+    language: str = typer.Option("es", help=f"The language of the text. Options: {LANGUAGES}"),
+    doc_type: str = typer.Option("sentence", 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."),
+    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}.",
+    ),
+    model_name: str = typer.Option(EMBEDDING_MODELS[0], 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 Lenses: Visualize text embeddings in 2D using colors to represent perplexity values.
+    """
+    logger.info("Loading embedding model...")
+    model = load_model(model_name)
+    dimensionality_reduction_function = (
+        partial(get_umap_embeddings, random_state=SEED) if dimensionality_reduction.lower() == "umap" else partial(get_tsne_embeddings, random_state=SEED)
+    )
+    logger.info("Loading KenLM model...")
+    kenlm_model = KenlmModel.from_pretrained(language)
+    logger.info("Loading dataset...")
+    if dataset.endswith(".csv") or dataset.endswith(".tsv"):
+        df = uploaded_file_to_dataframe(dataset)
+        if doc_type.lower() == "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(dataset, dataset_config, 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, seed=SEED)
+        logger.info("Saving plot")
+        bokeh_output_file(output_file)
+        save(plot)
+        logger.info("Done")
+
+
+if __name__ == "__main__":
+    app()

perplexity_lenses/data.py

@@ -17,7 +17,7 @@ def hub_dataset_to_dataframe(
     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":
+    if doc_type.lower() == "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])})

perplexity_lenses/engine.py

@@ -0,0 +1,110 @@
+import logging
+import time
+from typing import Callable, Optional, Union
+
+import pandas as pd
+import streamlit as st
+from bokeh.plotting import Figure
+from embedding_lenses.embedding import embed_text
+from embedding_lenses.utils import encode_labels
+from sentence_transformers import SentenceTransformer
+
+from perplexity_lenses.visualization import draw_interactive_scatter_plot
+
+logger = logging.getLogger(__name__)
+EMBEDDING_MODELS = ["distiluse-base-multilingual-cased-v1", "all-mpnet-base-v2", "flax-sentence-embeddings/all_datasets_v3_mpnet-base"]
+DIMENSIONALITY_REDUCTION_ALGORITHMS = ["UMAP", "t-SNE"]
+DOCUMENT_TYPES = ["Whole document", "Sentence"]
+SEED = 0
+LANGUAGES = [
+    "af",
+    "ar",
+    "az",
+    "be",
+    "bg",
+    "bn",
+    "ca",
+    "cs",
+    "da",
+    "de",
+    "el",
+    "en",
+    "es",
+    "et",
+    "fa",
+    "fi",
+    "fr",
+    "gu",
+    "he",
+    "hi",
+    "hr",
+    "hu",
+    "hy",
+    "id",
+    "is",
+    "it",
+    "ja",
+    "ka",
+    "kk",
+    "km",
+    "kn",
+    "ko",
+    "lt",
+    "lv",
+    "mk",
+    "ml",
+    "mn",
+    "mr",
+    "my",
+    "ne",
+    "nl",
+    "no",
+    "pl",
+    "pt",
+    "ro",
+    "ru",
+    "uk",
+    "zh",
+]
+
+
+class ContextLogger:
+    def __init__(self, text: str = ""):
+        self.text = text
+        self.start_time = time.time()
+
+    def __enter__(self):
+        logger.info(self.text)
+
+    def __exit__(self, type, value, traceback):
+        logger.info(f"Took: {time.time() - self.start_time:.4f} seconds")
+
+
+def generate_plot(
+    df: pd.DataFrame,
+    text_column: str,
+    label_column: str,
+    sample: Optional[int],
+    dimensionality_reduction_function: Callable,
+    model: SentenceTransformer,
+    seed: int = 0,
+    context_logger: Union[st.spinner, ContextLogger] = ContextLogger,
+) -> Figure:
+    if text_column not in df.columns:
+        raise ValueError(f"The specified column name doesn't exist. Columns available: {df.columns.values}")
+    if label_column not in df.columns:
+        df[label_column] = 0
+    df = df.dropna(subset=[text_column, label_column])
+    if sample:
+        df = df.sample(min(sample, df.shape[0]), random_state=seed)
+    with context_logger(text="Embedding text..."):
+        embeddings = embed_text(df[text_column].values.tolist(), model)
+    logger.info("Encoding labels")
+    encoded_labels = encode_labels(df[label_column])
+    with context_logger("Reducing dimensionality..."):
+        embeddings_2d = dimensionality_reduction_function(embeddings)
+    logger.info("Generating figure")
+    plot = draw_interactive_scatter_plot(
+        df[text_column].values, embeddings_2d[:, 0], embeddings_2d[:, 1], encoded_labels.values, df[label_column].values, text_column, label_column
+    )
+    return plot

requirements.txt

@@ -6,5 +6,6 @@ sentence-transformers==2.0.0
 bokeh==2.2.2
 umap-learn==0.5.2
 numpy==1.20.0
-https://files.pythonhosted.org/packages/f1/4a/ffebb6203694b9d2c9fdaafcdecba553ff039a0556804e88450c20ab9c73/embedding_lenses-0.4.0-py3-none-any.whl
-https://github.com/kpu/kenlm/archive/master.zip 
+https://files.pythonhosted.org/packages/2f/58/e00d2495b54f4ba97ca31a11aa7e636f80183ccf9b616f7eaa5518d050bb/embedding_lenses-0.5.0-py3-none-any.whl
+https://github.com/kpu/kenlm/archive/master.zip
+typer==0.4.0