Adjust batch size for extrating tokens

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:50c50c05859f43aa6a08aa3106a1ca62d225f1ac927d57e0e86e422cff5ee7a7
-size 711588089
+oid sha256:5ff31ebb2460dbc41a160cc755d0555bb8c84672563808b968a2a121c1b2414a
+size 711587941

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

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4265b625a915f8a622926c9be27d6b1f3f2bc44481f81ab5d53eace54a0bc06
+size 1421780139

tokens.py.orig

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+from datasets import load_dataset
+from tokenizers import ByteLevelBPETokenizer
+
+# Load dataset
+<<<<<<< HEAD
+dataset = load_dataset("oscar", "unshuffled_deduplicated_es", split="train[:5000000]")
+
+# Instantiate tokenizer
+tokenizer = ByteLevelBPETokenizer()
+def batch_iterator(batch_size=100_000):
+=======
+dataset = load_dataset("oscar", "unshuffled_deduplicated_es", split="train")
+
+# Instantiate tokenizer
+tokenizer = ByteLevelBPETokenizer()
+def batch_iterator(batch_size=1_000_000):
+>>>>>>> d5cede47e74aa6ec36f20acf5aba37c6734c6186
+    for i in range(0, len(dataset), batch_size):
+        yield dataset["text"][i: i + batch_size]
+
+# Customized training
+tokenizer.train_from_iterator(batch_iterator(), vocab_size=50265, min_frequency=2, special_tokens=[
+    "<s>",
+    "<pad>",
+    "</s>",
+    "<unk>",
+    "<mask>",
+])
+# Save files to disk
+tokenizer.save("./tokenizer.json")

tsne_plot.py

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