
What follows is research code. It is by no means optimized for speed, efficiency, or readability.

## Data loading, tokenizing and sharding

```python
import os
import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.decomposition import TruncatedSVD
from tqdm.notebook import tqdm
from openTSNE import TSNE
import datashader as ds
import colorcet as cc


from dask.distributed import Client
import dask.dataframe as dd
import dask_ml
import dask.bag as db

from transformers import AutoTokenizer
from datasets import load_dataset
from datasets.utils.py_utils import convert_file_size_to_int


def batch_tokenize(batch):
    return {'tokenized': [' '.join(e.tokens) for e in tokenizer(batch['text']).encodings]} # "text" column hard encoded

# The original viz used a subset of the ROOTS Corpus.
# More info on the entire dataset here: https://huggingface.co/bigscience-data
# And here: https://arxiv.org/abs/2303.03915

dset = load_dataset(..., split="train")

dset = dset.map(batch_tokenize, batched=True, batch_size=64, num_proc=28)

dset_name = "roots_subset"
max_shard_size = convert_file_size_to_int('300MB')
dataset_nbytes = dset.data.nbytes
num_shards = int(dataset_nbytes / max_shard_size) + 1
num_shards = max(num_shards, 1)
print(f"Sharding into {num_shards} files.")
os.makedirs(f"{dset_name}/tokenized", exist_ok=True)
for shard_index in tqdm(range(num_shards)):
    shard = dset.shard(num_shards=num_shards, index=shard_index, contiguous=True)
    shard.to_parquet(f"{dset_name}/tokenized/tokenized-{shard_index:03d}.parquet")

```

## Embedding

```python
client = Client() # To keep track of dask computation
client

df = dd.read_parquet(f'{dset_name}/tokenized/')
vect = dask_ml.feature_extraction.text.CountVectorizer(tokenizer=str.split,
                                                       token_pattern=None,
                                                       vocabulary=vocab)
tokenized_bag = df['tokenized'].to_bag()
X = vect.transform(tokenized_bag)

counts = X.compute()
client.shutdown()

tfidf_transformer = TfidfTransformer(sublinear_tf=True, norm="l2")
tfidf = tfidf_transformer.fit_transform(counts)

svd = TruncatedSVD(n_components=160)
X_svd = svd.fit_transform(tfidf)

tsne = TSNE(
    perplexity=30, # not sure what param setting resulted in the plot
    n_jobs=28,
    random_state=42,
    verbose=True,
)

tsne_embedding = tsne.fit(X)
```

## Plotting

```python
df = pd.DataFrame(data=tsne_embedding, columns=['x','y'])
agg = ds.Canvas(plot_height=600, plot_width=600).points(df, 'x', 'y')
img = ds.tf.shade(agg, cmap=cc.fire, how='eq_hist')
ds.tf.set_background(img, "black")
```

![ROOTS Dataset Scatterplot](./datashader.png)
license:apache-2.0 arxiv:2303.03915 region:us
{"license": "apache-2.0", "dataset_info": {"features": [{"name": "x", "dtype": "float64"}, {"name": "y", "dtype": "float64"}, {"name": "language", "dtype": "string"}, {"name": "corpus", "dtype": "string"}], "splits": [{"name": "train", "num_bytes": 247037602, "num_examples": 5785741}], "download_size": 112131877, "dataset_size": 247037602}}
TAGS
#license-apache-2.0 #arxiv-2303.03915 #region-us 


What follows is research code. It is by no means optimized for speed, efficiency, or readability.

## Data loading, tokenizing and sharding



## Embedding



## Plotting



!ROOTS Dataset Scatterplot
## Data loading, tokenizing and sharding ## Embedding ## Plotting



!ROOTS Dataset Scatterplot
TAGS
#license-apache-2.0 #arxiv-2303.03915 #region-us 
 ## Data loading, tokenizing and sharding ## Embedding ## Plotting



!ROOTS Dataset Scatterplot
passage: TAGS
#license-apache-2.0 #arxiv-2303.03915 #region-us 
## Data loading, tokenizing and sharding## Embedding## Plotting



!ROOTS Dataset Scatterplot