# Webdataset

WebDataset is a Python library that provides a convenient way to work with large datasets that are stored in a remote location, such as an Amazon S3 bucket or a Google Cloud Storage bucket. The library allows to stream data from these remote locations on the fly during the training.

To work on HFC, we need to use this type of solution because we don't have the space to store our data on disk.

## Setup
Incidentally, to stream data from s3 you need:
- install s5cmd with `conda install -c conda-forge s5cmd`
- configure aws with yours credentials
- put a specific command instead of the local paths that were used before, such as `pipe:bash ${WORKING_DIR}/experiments/pretraining/vloom/common/webdataset_get_file.sh ${tar_file_path_on_s3}`
- set in the configuration use_webdataset to true

## Philosophy of webdataset

The idea of Webdataset is to define a sequence of operations, each of which will be applied to the iterable resulting from the previous operation.

Here is an outline of the sequence of operations carried out:
- Input: a list or an iterable of commands or local paths. There are two identified use cases for m4. Either you want to open tar files stored locally, in which case the list should simply contain the local paths to the shards files. Or you want to open tar files stored on s3 and in this case you need to pass a command - `pipe:bash <PATH_TO_M4_CLONE>/experiments/pretraining/vloom/common/webdataset_get_file.sh <S3_URI>` - filled with the uri of the tar file.
- Sharding by node: Each gpu in the training will receive a subset of the iterable of input files. If this list has no duplicates, this means that each gpu will have unique shards.
- Sharding by worker: Each dataloader on each GPU can have several workers. If this is the case, the list of files assigned to a GPU will again be divided between the workers.
- Conversion of the tar file into samples: This step involves re-grouping the tar files that make up a single example. Webdataset is currently not designed for webdocuments that are made up of several texts and images, which is why we have a customised method `group_by_keys_interleaved`. This method also ensures that each example is complete, as piping can result in tar files being cut off in the middle.
- Collation of the samples into instances: This method involves changing the format of the samples to get closer to the format expected by our pipeline.
- Decoding of the images and the texts: So far, the elements in the instances are still bytes. This step converts them to their final python objects (PIL image and string)
- Batching of the examples: To finish the pipeline, we batch `map_batch_size` examples together

## Shuffling

To shuffle the examples, we can either:
1. Pseudo shuffle on the fly
2. Shuffle upstream

### Pseudo shuffle on the fly

When the tar files are not shuffled, or when several epochs are carried out on the same tar files, this is the method available to us.

The idea of on-the-fly pseudo-shuffling is to add buffers between targeted pipeline operations in order to randomly select the examples in this buffer. We have configuration variables to define the length of each buffer, but unfortunately the larger the buffer, the more RAM it consumes. In addition, webdataset offers to add a wrap-up phase during which you can start to draw from the buffer without having to wait until the buffer is completely full. By default, there is no shuffling on the fly, the following variables need to be adjusted for each dataset:
```python
shuffle_initial_urls_list: bool
shuffle_before_split_by_node_buffer_size: Optional[int]
shuffle_before_split_by_worker_buffer_size: Optional[int]
shuffle_after_tarfile_to_samples_buffer_size: Optional[int]
shuffle_after_batching_buffer_size: Optional[int]
```

### Shuffle upstream

The idea here is to store into the tar files the samples in a random order so that there is not pseudo-shuffling to do on the fly. In that case, the configuration parameters that need to be set to None/False are:
```python
shuffle_initial_urls_list: bool = False
shuffle_before_split_by_node_buffer_size: Optional[int] = None
shuffle_before_split_by_worker_buffer_size: Optional[int] = None
shuffle_after_tarfile_to_samples_buffer_size: Optional[int] = None
shuffle_after_batching_buffer_size: Optional[int] = None
```

## Resume training

Currently, we don't have a feature to resume a training where from where it left on the previous run in term of data (see "Potential improvements" section).

## Hyper-parameters tuning

On-the-fly streaming of examples from s3 adds hyper-parameters that have to be tuned for almost every scale of experiment. The parameters that will have the greatest influence on each other are: `max_num_images`, `max_seq_len`, `map_batch_size`, `max_num_samples_per_document`, `shuffle_before_split_by_node_buffer_size`, `shuffle_before_split_by_worker_buffer_size`, `shuffle_after_tarfile_to_samples_buffer_size`, `shuffle_after_batching_buffer_size`, `batch_size_per_gpu`, `num_workers` and the time of the forward + backward + opt step.

## Potential improvements

### S3 pipe script

Currently, the s3 piping works by downloading the shard to the node's NVME drive and then piping the file. This solution appears to be sub-optimal because there is an extra write and read on the NVME. However, without this trick, we don't have the full tar files in the pipe, we never get to the end. The current hypothesis is that the internal retry system of libs such as `s5cmd` or `aws s3` do not work with the pipe.

### Disallow list

To have good control over the non-repetition of data in a training that is split into several jobs, a disallow list system should be implemented and used in conjunction with upstream shuffling. It's not a simple feature, especially if you want a perfect implementation. Nevertheless, if you accept losing the end of a few shards, the solution shown in the PR #1307 should provide a good basis.

## Create tar files

All the files used to create the tar files are inside `datasets_processing_scripts/01_tar_datasets_with_jpeg`. For future processing, particular attention should be paid to the order of the files in the tar.

## Debug tips

Currently, if there is a bug in the webdataset pipeline, it will not cause the training to crash. The error will simply be logged and the code will move on. For the future or for debugging, the following change should be considered: `handler=log_and_continue` -> `handler=wds.reraise_exception`.

# Checkpoint extraction

At the end of the training the normal model weights file isn't in the checkpoint and requires a manual extraction, which is done offline and the script now has the luxury of using the whole node's CPU RAM, e.g.

```

cd /fsx/m4/experiments/local_experiment_dir/tr_171-save-load/opt_step-50/accelerator_state/
./zero_to_fp32.py . output.bin
```

The `zero_to_fp32.py` script is already copied into the checkpoint upon checkpoint saving.

We aren't gathering the full model on every save because it's slow and there might not be enough memory to perform that. Therefore we use `stage3_gather_16bit_weights_on_model_save: False` to only having each gpu save its own shards.

# Monitoring a training

What does it mean to monitor a training? The most important things to look at:
- Most importantly, the loss should go down (on average) consistenly. If it diverges, intervention is required (see next section)
- Looking at the colab metrics (with parameters/gradients/activations) is a useful indicator. Weird behaviors (as in explosion) usually precedes a divergence of the loss.
- Is the training still in the slurm queue? If not, intervention is required (see next section)
- Nodes failures (it will most likely make the training crash). You can do a `grep 'srun: error: Node failure'` on the text logs. Reporting them on #science-cluster-support is a good idea.

# How to intervene when a training diverges

In case of rewinding, I recommend starting another plot on WB by setting the `wandb_run_id` inside `resume_run_infos.json` to the empty string. This will create a new run on WB and you can compare the two runs, and hopefully the loss for that new run will not diverge.

Try in order of complexity:
- Rewind and restart (that is essentially reshuffling the data)
- Rewind, decrease the LR and restart
- Rewind, reset the optimizer, set lr=0, restart, train with lr=0 for a bit, then restart again now with restored lr

## How to reset the wandb run on rollback

empty the `wandb_run_id` string in the `resume_run_infos.json`

## How do I rewind to a previous training?

Change the path of the resuming checkpoint in `$SAVE_DIR/latest_opt_step_dir`.

Additionally, if you already started to evaluate the previous checkpoints that you are now discarding, you might want to backup your previous results and roll back your evaluations results json file to the same opt step.

Here's an example script to do it:
```python

from pathlib import Path
import json
import shutil

# ------------ Fill in these variables ------------
roll_back_number = XX # e.g. 4
roll_back_to_opt_step = XX  # e.g. 16000
run_name = "XXX" # e.g. "tr_190_01_64n"
# -------------------------------------------------

main_eval_path = Path(f"/fsx/m4/experiments/local_experiment_dir/evals/results/{run_name}_evaluations.jsonl")
archive_eval_path = main_eval_path.parent / f"{run_name}_evaluations_archive_exp_{roll_back_number}.jsonl"

# First we start by backing up the current evals file
shutil.move(main_eval_path, archive_eval_path)

# Then we select the evals we want to keep

# 1. Load the evals from the archive
parsed_json_objects = []
try:
    with open(archive_eval_path, 'r', encoding="ISO-8859-1") as file:
        json_data = file.read()

    # Split the JSON data into separate entities
    json_entities = json_data.strip().split('\n')

    for entity in json_entities:
        try:
            parsed_json = json.loads(entity)
            parsed_json_objects.append(parsed_json)
        except json.JSONDecodeError as e:
            print("Failed to parse JSON entity:", e)

    print("JSON entities parsing succeeded.")
except IOError as e:
    print("Error reading the file:", e)

# 2. Select the evals we want to keep
cleaned_parsed_json_objects = []
for parsed_json in parsed_json_objects:
    curr_opt_step = int(parsed_json["model_name_or_path"].split("/")[-2].split("-")[-1])
    if curr_opt_step <= roll_back_to_opt_step:
        cleaned_parsed_json_objects.append(parsed_json)

print(len(cleaned_parsed_json_objects), len(parsed_json_objects))

with open(main_eval_path, 'w') as output_file:
    for r in cleaned_parsed_json_objects:
        output_file.write(json.dumps(r) + "\n")
    print("Parsed JSON data saved to", main_eval_path)

```

## How do I decrease/change the LR mid-training?

Tweak the states (there are two fields for the LR) in `opt_step-xxx/accelerator_state/custom_checkpoint.pkl`.

e.g. to reduce lr by 25% do:
```
python -c "import sys, torch; sd=torch.load(sys.argv[1]); \
print(sd['base_lrs']); sd['base_lrs'] = [x*0.75 for x in sd['base_lrs']]; print(sd['base_lrs']); \
print(sd['_last_lr']); sd['_last_lr'] = [x*0.75 for x in sd['_last_lr']]; print(sd['_last_lr']); \
torch.save(sd, sys.argv[1])" opt_step-14500/accelerator_state/custom_checkpoint_0.pkl
```

## How do I reset the optimizer?

Set `load_optimizer_states` to False.

# I detected a bug in the code, and it will require some time to fix.

If you think that fixing a bug might take more than 1 day and fixing it doesn't require the full capacity (in terms of node) that we are training on, then put the job array on hold (`scontrol hold <job_id>`), let people know #science-cluster-planning that they should use the idle GPUs with some other runs. Right now bigcode almost have some smaller jobs to run on the side, so they can squeeze in some jobs and halt them when we are ready to relaunch.


# Generation Process:

- find one or more opt-step checkpoints to make generations with
- create folder in code/m4/experiments/generations
- add a config.yaml and a [gen_folder_name]_generate.slurm folder
- fill the config file according to desired hyperparameters: prompt/num_beams/ngram_repeats etc..
- run sbatch [m4_repo_name]/experiments/generation/[gen_folder_name]/[gen_folder_name]_generate.slurm
- check wandb and make sure your column shows up. If it doesn't, click on "columns" at the bottom right of the generation table and slide the missing generation to the "Displayed columns" side

Everything about the SmolLM2 and SmolVLM family of models


We need to locally save some datasets with `copy_remote_sample_datasets.py` because the caching function does not work for some datasets, see https://github.com/huggingface/datasets/issues/4760 and https://github.com/huggingface/datasets/issues/3547.


# Data Collection

## Goal of `data_collection`

This folder is aimed to:
- Simplify HTML DOM trees;
- Convert the simplified DOM trees to another structure adapted for an extraction;
- Perform an extraction (either of image-text pairs, or web documents);
- Perform a filtering on the extraction (either on image-text pairs, or web documents);
- Visualize the results.

## Organization of `data_collection`

The sub-folder `processors` contains the files defining the functions to do the big operations:
- The simplification of DOM trees in `dom_tree_simplificator.py`;
- The convertion of DOM trees to a more adapted structure in `pre_extraction_simplificator`;
- The extraction of web documents in `web_document_extractor.py`;
- The filtering of web documents in `web_document_filtering.py`;
- The extraction of pairs in `pair_extractor.py`;
- The filtering of pairs in `pair_filtering.py`.

These files require other functions or external parameters to work, which are defined in the sub-folder `utils`.

The call of these operations from `processors` to obtain outputs (in `outputs`, which is now a bit outdated, or rather `large_files` at the root of the repo) are done in the sub-folder `callers`.

The sub-folder `visualization` contains the streamlit apps to visualize these outputs:
- In `global_visualization.py`, one can view how the simplification of DOM trees affect the HTML codes, the trees, and the rendered webpages. We also visualize web documents and extract image-text pairs with additional information.
- In `web_document_visualization.py`, one can visualize web documents and see the impact of filtering on them.
- In `pair_visualization.py`, one can obtain statistics on the extracted image-text pairs, and see the impact of filtering on these statistics.
- `plot_clib_distrib.py` is used to obtain the distributions of CLIP scores of reference datasets, to compare with our distribution.
- We used `pair_stat_dashboard.py` to obtain a lot of statistics on the pairs at the beginning. This file might not be maintained anymore.

The sub-folder `debug` is used for debugging and could be ignored.

## Explanation of the data collection pipeline

### Starting point for the dataset

We start with the dataset [`c4-en-html-with-metadata`](https://huggingface.co/datasets/bs-modeling-metadata/c4-en-html-with-metadata), which contains 45M of English HTML documents whose url corresponds to C4 examples gathered by the modeling metadata group.

Each example of the dataset contains much information about the metadata, but we are currently only interested in the columns `html` and `url`.

The full dataset is downloaded on Jean Zay at the path `/gpfsscratch/rech/cnw/commun/local_datasets/c4-en-html-with-metadata-arrow/` (5.9T), and the full dataset containing only the columns `html` and `url` is at the path `/gpfsscratch/rech/cnw/urd43gx/c4_en_html` (3.5T).

### From an HTML string to a tree structure

We use [Selectolax](https://github.com/rushter/selectolax) to efficiently parse the HTML strings and create trees. Each node in the tree corresponds to a tag or a text, and we can have information about its attributes.

### Simplifying the DOM trees

With the `processor` `DOMTreeSimplificator`, we make several simplifications to have simplified DOM trees and remove the unnecessary information:
- We remove comments in the HTML strings;
- We replace the tags `<br>`, `<br/>` and `<br />` by a break of line;
- We unwrap a list of tags (see below), meaning that we remove the tags, but we keep the content inside;
- We strip every tag not in a list (see below), meaning that we completely remove both the tags and what’s inside;
- In a complementary way to the previous step, we additionally remove some tags which contains wrong attribute values, for example a `<div>` tag containing an attribute `class` with the value `date`;
- We remove empty nodes;
- We un-nest nodes, meaning that if a parent node only has one child and no text associated, the child becomes the parent.

Tags that we unwrap: `a`, `abbr`, `acronym`, `b`, `bdi`, `bdo`, `big`, `cite`, `code`, `data`, `dfn`, `em`, `font`, `i`, `ins`, `kbd`, `mark`, `q`, `s`, `samp`, `shadow`, `small`, `span`, `strike`, `strong`, `sub`, `sup`, `time`, `tt`, `u`, `var`, `wbr`.

After unwrapping the tags, tags **not** in this list are removed:
- Tags defining a structure : `address`, `article`, `aside`, `blink`, `blockquote`, `body`, `br`, `caption`, `center`, `dd`, `dl`, `dt`, `div`, `figcaption`, `h`, `h1`, `h2`, `h3`, `h4`, `h5`, `h6`, `hgroup`, `html`, `legend`, `main`, `marquee`, `ol`, `p`, `section`, `summary`, `title`, `ul`;
- Tags defining a media: `audio`, `embed`, `figure`, `iframe`, `img`, `object`, `picture`, `video`;
- Tags that could contain an interesting attribute: `source`.

Tags that we could consider if we really want to have a high recall (but we cannot do anything with most of them): `bgsound`, `button`, `canvas`, `col`, `colgroup`, `datalist`, `details`, `dialog`, `dir`, `fieldset`, `form`, `frame`, `frameset`, `head`, `header`, `input`, `li`, `label`, `map`, `nav`, `optgroup`, `option`, `pre`, `select`, `slot`, `svg`, `table`, `tbody`, `td`, `template`, `tfoot`, `th`, `thead`, `tr`, `track`, `xmp`.

For the 3rd point, they are tags that we currently chose to strip, but that we might reconsider later: `table` (and its associated tags), `form`, `li`, `head`, `header`, `nav`. We chose to remove these tags, either because it is hard to transform to a text (how do we transform a `table` to something clear with a linear text?), or because it can contain useful information, but in most cases this is noise of information related to the navigation in the website (`li` for example).

### Having a good structure

With the `processor` `PreExtractionSimplification`, we traverse the simplified DOM tree and append the nodes to a list to have a better structure.

If the node is a media node, we extract the interesting attributes. We check the validity of the source URL, if it’s not valid we discard the media node.

If the node is a text node, then we format it with the following strategy:
- In the HTML string, replace every +2 `\n` with only one `\n`;
- In the HTML string, replace every +2 spaces with only one space;
- In the HTML string, replace `<br> tags` (and their various forms) with `#BR_TAG#`;
- Within a text node, replace every `\n` or `\t` with a space;
- Within a text node, replace every +2 spaces with only one space;
- Within a text node, split on `#BR_TAG#`, strip each element, and merge on `\n`. If the very first and/or last characters of the text are spaces, make sure to keep them.

Then, we have the possibility to merge two consecutive text nodes (and repeat this operation) with the following strategy:
- Append the separation at the end of the first text (if there is any, it will be one space) to a set, and remove it from this text;
- Append the separation at the beginning of the second text (if there is any, it will be one space) to the set, and remove it from this text;
- Consider all the tags that differ from the path of node 1 to the path of node 2. Append to the set the separations induced by each of these tags;
- The biggest separation in the set wins, in this order: `\n\n`, `\n`, space, nothing. Merge the two texts with this separation;
- When a text node cannot be merged anymore (the previous and following nodes are not text nodes), strip the text.

### Intuition behind Selectolax trees and merging text nodes

Selectolax builds trees by creating a new node for each tag present in the HTML document, and a new node for each non-empty text.

For exemple, consider the HTML document:
```
html_str = """
<html>
<body>
<div>
    this is a test
    <h1>Heading</h1>
    <p>
        p-inner
    </p>
    p-trailing
</div>
</body>
</html>
"""
```

When traversing the tree (depth-first) and printing the path of the nodes (which include their tag as the last component of the path), we obtain:
- `.html`
- `.html.head`
- `.html.body`
- `.html.body.-text` (the text is "\n", because there is a line break between `body` and `div`)
- `.html.body.div`
- - `.html.body.div.-text` (the text is "\n\tthis is a test\n\t")
`.html.body.div.h1`
- `.html.body.div.h1.-text` (the text is "Heading")
- `.html.body.div.-text` (the text is "\n\t")
- `.html.body.div.p`
- `.html.body.div.p.-text` (the text is "\n\t\tp-inner\n\t")
- `.html.body.div.-text` (the text is "\n\tp-trailing\n")
- `.html.body.-text` (the text is "\n\n")

Now we want to merge these text nodes.

We first merge the first text node, "\n" at `.html.body.-text`, and the second text node, "\n\tthis is a test\n\t" at `.html.body.div.-text`.

To do that, we follow the strategy by first formatting the two texts, which results in “ “ and " this is a test ".

We notice that in the first text “ “, we have a space separation at the beginning of the text (it is a particular case since it is also at the end here, but considering it at the end would not change the result in the end). So we'll need to keep this space at the beginning of the merged text nodes, and we can remove it from the text 1, which becomes the empty text “”.

We notice that there isn’t any separation at the end of the text 1, but there is a separation at the beginning of the text 2, which is a space. So we add this space to our set of separations, and we remove it from the text 2, which becomes “this is a test “.

Now, we simply have to check the differences between the paths of the two text nodes. We only have `div` which is different, and `div` induces a line break `\n`, so we add “\n” to our set of separation.

Our set of separations includes “ “ and “\n”. The strongest is “\n”, so this will be our separation between the two text nodes, which becomes “ “ (the first leading space that we should not forget about) + “” (text 1) + “\n” (separation) + “this is a test “ (text 2) = “ this is a test “.

This merged text nodes takes the path of the second text nodes, which is `.html.body.div.-text`.

We can now merge this new merged text node “ this is a test “ at .`html.body.div.-text` with “Heading” at `.html.body.div.h1.-text`. This results after the operation in “ this is a test\n\nHeading” at `.html.body.div.h1.-text`.

And so on until merging everything, which results in “ this is a test\n\nHeading\n\np-inner\n\np-trailing\n\n”.

Since we cannot merge this node anymore, we can strip it to obtain “this is a test\n\nHeading\n\np-inner\n\np-trailing”. This is what is rendering by testing online on an [HTML editor](https://htmledit.squarefree.com/).

### Web documents extraction

At this stage, web documents are simply the structure detailed previously, where each node is either a text or an image (or nothing if we couldn’t download the image).

We use the the `processor` `CommonCrawlWebDocumentExtractor` to extract web documents. Before doing this, make sure to [set up a DNS resolver](https://github.com/rom1504/img2dataset#setting-up-a-bind9-resolver).

Performances for 10K documents, on a Mac M1 Pro (all steps are done with multiprocessing with 10 cores) **without** resizing images:
- Step 1 - Extracting and processing the HTML files: 15 secs
- Step 2 - Getting the URLs of all images: < 1 sec
- Step 3 - Downloading the images with `img2dataset`: 4 min
- Step 4 - Create the dataset containing all images (**2 Go**): 7 secs
- Step 5 - Replace the URLs in the dataset obtained after Step 1 with images bytes (image retrieval): 2 secs

On 1M documents on a GCP machine (60 cores):
- Processing the HTML documents, simplifying them, and creating a dataset with the desired structure: 7min
- Download images: 1h24min
- Create the dataset containing all images: 30min
- Retrieve images and add them to the initial dataset: 16min

1964081 out of 4034154 images were successfully downloaded (48.7%).
The dataset containing all images weighs 185G.
The final dataset with texts and images weighs 243G.

On 10M documents on a GCP machine (60 cores):
- Processing the HTML documents, simplifying them, and creating a dataset with the desired structure: 1h03min
- Download images: 11h15min
- Create the dataset containing all images: 6h29min
- Retrieve images and add them to the initial dataset: 2h25min
- Saving the dataset : 1h36min
- Sharding the dataset : 7h29min

### Web document filtering

The filtering of web documents is done at different levels. First, we modify the document with a filtering at node level. Then, we decide if we keep the document with a filtering at document level.

**Node level:**

For each image, filters on:
- The format;
- The size (original and rendered widths and heights, side ratio).

For each paragraph in a text, filters on:
- The number of words;
- The special character ratio;
- The stop word ratio.

**Doc level:**

Filters on:
- The number of images;
- The number of words;
- The character repetition ratio;
- The word repetition ratio;
- The special character ratio;
- The stop word ratio;
- The flagged word ratio;
- The language identification prediction score;
- The perplexity score.

### Image-text pairs extraction

With the `processor` `TextMediaPairsExtractor`, we extract image-text pairs first by looking at images in the list of nodes from the previously described structure. We only keep images that we are able to download. Then, to form pairs with an image, we consider the alt-text (if present), we consider the formatted filename (essentially applying some regexes to prettify the original name of the file), and we consider the extracted text if the node just after the image is a text node. We then split this text on `\n\n` and consider the first element (essentially the first paragraph). We also have the possibility to extract images not present in the simplified DOM tree, but in this case the extracted text is never present.

### Image-text pairs filtering

With the `processor` `PairFiltering`, we essentially filter image-text pairs based on:
- The format of the images;
- The size of the images (original and displayed width, original and displayed height, side ratio);
- The number of words in the texts;
- The special character ratio in the texts;
- The repetition ratio in the texts;
- The CLIP scores.


# Intersect a list of urls with the urls archived in a snapshot of Common Crawl

In this section, I want to leave a trace of the steps I followed in order to know how many oscar's english subset urls are present in the CC-MAIN-2021-25 snapshot of Common Crawl

## 1. Get the list of urls from oscar's english subset

```python
from datasets import load_dataset
import json
from tqdm import tqdm

saving_path = "/gpfswork/rech/cnw/urd43gx/urls_oscar_english/urls_oscar_english.parquet" #CHANGEME

dataset = load_dataset("oscar-corpus/OSCAR-2109", language="en", split="train", use_auth_token=True)

print("Dataset successfully loaded")

def get_urls_from_meta_column(meta_col):
   urls = [meta_item["headers"]["warc-target-uri"] for meta_item in meta_col]
   return {"url": urls}

dataset = dataset.map(get_urls_from_meta_column, batched=True, batch_size=1000, remove_columns=dataset.column_names, num_proc=25, input_columns=["meta"])

dataset.to_parquet(saving_path)
```
Note, for the following steps we need the list in a table in parquet format

## 2. Transfer the parquet table to a S3 bucket

I copied it here: `s3://m4-cc-index/urls_oscar_english/urls_oscar_english.parquet`

## 3. Create on [AWS Athena](https://aws.amazon.com/athena/) a database and a table storing the Common Crawl index

Follow the steps described here: https://commoncrawl.org/2018/03/index-to-warc-files-and-urls-in-columnar-format/

If the table `ccindex` already exist, don't forget to update it with the latest crawl snapshots by running `MSCK REPAIR TABLE ccindex`.

## 4. Create a new database and table with the oscar's english subset urls

On Athena UI, run:

```sql
CREATE DATABASE m4
```

```sql
CREATE EXTERNAL TABLE IF NOT EXISTS m4.urls_oscar_english (
         `url` string)
ROW FORMAT SERDE 'org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe'
WITH SERDEPROPERTIES (
  'serialization.format' = '1'
) LOCATION 's3://m4-cc-index/urls_oscar_english'
TBLPROPERTIES ('has_encrypted_data'='false');
```

```sql
SELECT * FROM "m4"."urls_oscar_english" limit 10;
```

|url                                                                                |
|--------------------------------------------------------------------------------------------------------------------------------------------------------------|
|https://cupe4914.com/cupe-4914-warns-that-peel-cas-needs-to-stop-putting-vital-supports-at-risk-and-get-serious-about-negotiating-a-fair-collective-agreement/|
|https://cure-ed-info.com/politics/shameful-publicity-stunt-labour-mp-hits-out-at-rishi-sunak-sit-down-with-gordon-ramsay/                                     |
|https://cure-ed-info.com/world-news/tanzanian-president-blames-lab-after-goat-papaya-test-positive-for-coronavirus/                                           |
|https://custom-essay-writer.net/2020/10/21/literature-review-hq_7h/                                                                                |
|https://customclassicsmancaves.com/index.php?route=product/product&product_id=746                                                                             |
|https://customdistributors.com/recipe/blt-baked-potatoes/                                                                                |
|https://customtwit.com/30-laws-of-flow-wordpress-site-with-charlene-day/                                                                                |
|https://customwritingshelp.com/2021/02/23/question-1-of-20-one-purpose-of-closing-entries-is-to-give-zero/                                                    |
|https://customwritingwebsite.com/2020/06/30/write-my-college-essay_kl/                                                                                |
|https://cuttingedgechicago.com/terms-and-conditions                                                                                |

## 5. Join the  oscar's english subset urls table with the Common Crawl index table

```sql
CREATE TABLE "m4"."result_urls_oscar_english" WITH (
  format = 'parquet',
  external_location = 's3://m4-cc-index/result_urls_oscar_english/'
) AS
select a.*,
  b.*
from (
    select url
    from m4.urls_oscar_english
  ) as a
  left join (
    select url as url_2,
      url_host_name,
      content_mime_type,
      content_mime_detected,
      warc_filename,
      warc_record_offset,
      warc_record_length,
      warc_segment,
      crawl,
      fetch_status,
      content_languages
    from ccindex.ccindex
    where crawl = 'CC-MAIN-2021-25'
  ) as b on a.url = b.url_2
```

## 6. Get the number of oscar's english subset urls in the CC-MAIN-2021-25 snapshot
```sql
select count(*)
from m4."result_urls_oscar_english"
where url_2 is not NULL;
```
108551545

without duplicated urls:
```sql
select count(DISTINCT url)
from m4."result_urls_oscar_english"
where url_2 is not NULL;
```
106503003


# Data Processing Pipelines

Relate to issue [#12](https://github.com/huggingface/m4/issues/12).

We have two v0 data processing pipelines:
- (a) split (for sharding) + parallel/slurm arrays of whatever processing scripts (python or rust for instance)
- (b) apache beam (for creating processing pipelines) + Dataflow (for horizontal scaling)

## App

ngram search is mostly an example.
to launch the app:
```bash
streamlit run app.py --server.port 6006
```


# Web document filtering documentation

The filtering is done at node and document levels.

At node level, we consider paragraphs or images. This filtering is done to clean the document before doing the filtering at document level, which is deciding if we want to keep the document or not.

Some filters are defined at both node and document levels. If the thresholds were the same for these two levels, it wouldn't be useful to call these filters at document level again, since it would automatically pass the filtering, given the fact that we removed the problematic nodes at node level.

However, the thresholds shouldn't be the same at node and document levels. In a text, at node level, you can have short sentences while at document level, you see the bigger picture. Then, you can be much stricter on the threshold at document level while keeping a high recall than at node level.

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/utils/filtering_utils.py#L37

## Filtering at node level

This operation is done in a `.map`. So we have web documents as inputs, and modified web documents as outputs.

### Filtering at node levels for texts

We start by **modifying** the texts by doing:
- Non printing characters removal (just removing weird characters that are not rendered when displaying a text, but still visible for a computer);
- White space standardization (there are many different white space characters, and we need to standardize them to be able to split on white space characters to obtain the words in a text, useful later for the filtering).

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L61

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L65

Then, for each text, we split on `\n\n` to obtain paragraphs, and for each paragraph, we apply a filtering on:
- The number of words;
- The character repetition ratio;
- The word repetition ratio;
- The special character ratio;
- The stop word ratio;
- The flagged word ratio;
- The ratio of punctuation characters vs number of words;
- The common word ratio;
- The language identification prediction score;
- The perplexity score.

See details below for the calculations of these quantities.

We remove paragraphs that do not pass the filtering, and join the remaining paragraphs on `\n\n`.

### Filtering at node levels for images

For each image, we filter on:
- The format;
- The size (original and rendered widths and heights, aspect ratio).

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L13

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L20

## Filtering at document level

This operation is done in a `.filter`. We have the modified web documents as inputs (the outputs of the filtering at node level), and the booleans indicating if we keep the documents or not as outputs.

We filter web documents on:
- The number of images;
- The number of words;
- The character repetition ratio;
- The word repetition ratio;
- The special character ratio;
- The stop word ratio;
- The flagged word ratio;
- The ratio of punctuation characters vs number of words;
- The common word ratio;
- The language identification prediction score;
- The perplexity score.

See details below for the calculations of these quantities.

## Keeping a high diversity

Some web documents, even if they are rare, could be extremely long, or contain an inordinate number of images.

Imagine we have a dataset of 1000 documents, 999 of these documents contain 10 words and 1 image each, and the remaining document contains 10000 words and 1000 images.

Then, we have the feeling that we have diversity in our dataset, since we took 1000 random documents from the internet, but half of the content is from one single document, which is likely to be on the same topic (if it is not spam, which is highly possible).

To remove these outliers (that exist, and take all the place in the dataset), we remove documents with too many words or images.

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L51

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L137

## Character repetition ratio calculation

For a given $n$, we count the occurrence of each *character* $n$-gram present in the text. We denote $r$ the number of character $n$-grams with at least 2 occurences. We define the character repetition ratio as the ratio of the sum of the $min(k, r)$ largest occurrences ($k$ is defined just below) by the sum of all occurrences, and we discarded texts with a too high ratio.

If $k=1$, short sentences are much more likely to have a high character repetition ratio, since the most frequent character $n$-gram represents a larger proportion of the sentence. If $k$ is the number of occurrences greater than or equal to $2$, very long texts, but not necessarily including repetitions, tend to have a high character repetition ratio, since these texts inherently have a wide diversity of character $n$-grams. $k=\lfloor \sqrt{N} \rfloor $, with $N$ the number of different character $n$-grams found in the text, counterbalances well this effect in practice.

*Example:* Take the sentence `ok_ok_good_ok` and $n=3$. Character $n$-grams, with their frequencies, are given in the following table.

| `ok_` | `_ok` | `k\_o` | `k\_g` | `\_go` | `goo` | `ood` | `od\_` | `d\_o` |
| - | - | - | - | - | - | - | - | - |
| 2 | 2 | 1 | 1 | 1 | 1 | 1 | 1 | 1 |

Since we have 9 different character $n$-grams, $N=9$ and $k = \lfloor \sqrt{N} \rfloor =3$.

We have two character $n$-grams with at least two occurrences, so $r=2$.

Then, $min(k, r)=2$.

The sum of the $min(k, r)$ largest occurrences is $2+2=4$ and the sum of all occurrences is $11$. Thus, the character repetition ratio for this sentence is $\frac{4}{11}$.

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L167

## Word repetition ratio calculation

As a complement to the previous filter, we remove texts that have commonly repeated similar long sentences. More specifically, we create a filter for the repetitions by looking this time at the occurrences of the *word* $n$-grams, for a chosen $n$ parameter. We define the word repetition ratio as the ratio of the sum of the occurrences greater than or equal to 2 to the sum of all occurrences, and we discard texts with too high of a ratio. Contrary to the filter on the character repetition ratios, I did not find a bias of this method giving systematically higher or lower scores to longer or short texts. This filter is more robust in finding texts with long exact duplicated sentences in them, while the previous one is used to find short to medium sized repetitions.

*Example:* Take the sentence `My name is Hugo. What is your name? My name is Paul.` and $n=2$. Word $n$-grams, with their frequencies, are given in the following table.

| `My name` | `name is` | `is Hugo` | `Hugo What` | `What is` | `is your` | `your name` | `name My` | `is Paul` |
| - | - | - | - | - | - | - | - | - |
| 2 | 2 | 1 | 1 | 1 | 1 | 1 | 1 | 1 |

We have two word $n$-grams with at least two occurrences, for a total number of $2+2=4$ occurences.
The sum of all occurrences is $11$, so the word repetition ratio for this sentence is $\frac{4}{11}$

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L197

## Special character ratio calculation

The list of special characters was defined by taking an existing list of special characters, then finding in many web texts non ASCII characters that were not present in this list and count their occurences, and finally adding the most frequent ones to the original list. Emojis are also added to the list. We simply discard texts with a special character ratio above a certain threshold.

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L217

## Stop word ratio calculation

Having a low stop word (or closed class word) ratio in a text is one of the best indicators of a non-human generated content. The list of stop words was built by taking pre-existing lists, for example from Universal Dependencies. We discard texts with a too low closed class word ratio.

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L236

## Flagged word ratio calculation

To build a list of flagged words, we used a circular method of 1 step:
- We started with a concatenation of pre-existing (even if not perfect) lists of flagged words found on the internet;
- Then, we computed the flagged word ratios of many documents from the internet, using this list;
- We used these scores to build a database containing the documents with the highest flagged words ratios;
- Then, we manually inspected these documents to discover new words to add to the list;
- Finally, the list was filtered with precise instructions below.

*Instructions for building the lists of flagged words:* Keep only the words associated with porn and systematically used in a sexual context. Remove words that can be used in medical, scientific, colloquial (without referring systematically to porn), or everyday contexts. Remove all insults. Remove all words referring to race or sexual orientation.

We are then able to compute the flagged word ratio of a text and discard it if it is too high.

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L255

## Punctuation ratio calculation

With a regular expression, we split a text string into a list of words and punctuations (predefined list of punctuation characters used in English). We then count the ratio between number of punctuations and number of words. We discard texts with too low punctuation ratio, as they are usually indicative of poor quality text.

## Common word ratio calculation

We analyze a large amount of text from the Common Crawl using the Oscar dataset, extracting and counting words and removing those that occur only once. We calculate the common word ratio of a text to identify machine-generated content, removing texts with a low ratio.

https://github.com/huggingface/m4/blob/57bda9f70eec539401046b5127ecdff5ae6b4e71/m4/sourcing/data_collection/processors/web_document_filtering.py#L317

## Language identification prediction score calculation

FastText is used to perform language identification and getting confidence scores for a text. If a score is below a specific threshold, we discard the text.

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L279

## Perplexity score calculation

SentencePiece unigram tokenizers followed by KenLM 5-gram models after tokenization were trained on Wikipedia article openings. We discarded texts with too high perplexity scores.

https://github.com/huggingface/m4/blob/4e95b234c1206355848faf0e77492717e3e70154/m4/sourcing/data_collection/processors/web_document_filtering.py#L363


# Image deduplication

# Methods

The main information of an image is contained in its low frequencies, the high frequencies providing detail. The main methods to perform image deduplication in a fast way, and whose quality is not so bad, are based on this observation.

## Average hashing (AHash)

- Convert the image to grayscale.
- Reduce the size of the image. To have a hash of 64, shrink the image to 8x8. It will remove the high frequencies.
- Compute the average value of the 64 pixels.
- For each pixel, replace its value by 1 if it is larger than the average or by 0 otherwise.
- Unroll the image to obtain the hash.

A hamming distance is used to compare two hashes. It is fast, but not efficient against minor modifications of the image. Lot of false positives.

## Perceptual hashing (PHash)

- Convert the image to grayscale.
- Reduce the size of the image, to 32x32 for example. This step is done to simplify the DCT computation and not because it is needed to reduce the high frequencies.
- Compute the 32x32 DCT, and keep the top left 8x8, which represents the low frequencies.
- Compute the average value of the top left 8x8, and exclude the DC coefficient.
- For each of the 64 pixels, replace its value by 1 if it is larger than the average or by 0 otherwise.
- Unroll the image to obtain the hash.

It is slower than AHash, but more robust to minor modifications of the image. Less false positives than AHash in practice.

## Difference hashing (DHash)

- Convert the image to grayscale.
- Reduce the size to 9x8, essentially to remove the high frequencies.
- For each pixel, compare its value to the one at the right. Replace by 1 if it is larger or by 0 otherwise, until ending with a 8x8 image.
- Unroll the image to obtain the hash.
- Optional: repeat the steps by comparing pixels on the columns instead of on the rows, and concatenate the new hash with the previous one to obtain a 128-bit hash, and in practice reduce the the number of false positives.

It is as fast as AHash, with less false positives. Less accuracy than PHash.

## Wavelet hashing (WHash)
Same as PHash, but uses DWT instead of DCT. It is way faster than PHash, a bit less than AHash and DHash, but produces way more false positives.

# Libraries

## [`imagededup`](https://github.com/idealo/imagededup)

It duplicates with the algorithms: CNN, Perceptual hashing, Difference hashing, Wavelet hashing, Average hashing.

## [`imagehash`](https://github.com/JohannesBuchner/imagehash)

It supports: Perceptual hashing, Difference hashing, Wavelet hashing, Average hashing, HSV color hashing (colorhash), Crop-resistant hashing.

## [`image-match`](https://github.com/ProvenanceLabs/image-match)

It implements the Image signature algorithm.

## [`imgdupes`](https://github.com/knjcode/imgdupes)

It supports Perceptual hashing (using only the 8x8 DCT low-frequency values including the first term), Difference hashing, Wavelet hashing, Average hashing, Perceptual hashing org (using only the 8x8 DCT low-frequency values and excluding the first term since the DC coefficient can be significantly different from the other values and will throw off the average). It uses `imagehash` except for Perceptual hashing org.

## [`simhash-py`](https://github.com/seomoz/simhash-py)

Implements the SimHash algorithm, as well as a solution for the Hamming distance problem.

## [`faiss`](https://github.com/facebookresearch/faiss)

Efficient similarity search and clustering of dense vectors.

## [`hnswlib`](https://github.com/nmslib/hnswlib)

Fast approximate nearest neighbors

# Research papers

## Interesting papers

- [Duplicate Discovery on 2 Billion Internet Images (2013)](https://people.csail.mit.edu/celiu/pdfs/CVPR13-bigvis-dupdiscovery.pdf):
This paper presents a method to provide a hash to an image, by considering different scales and splitting the image, computing the average pixel values of each block, gathering everything, performing a PCA (trained on a portion of the total of the images), and quantizing it by putting to 0 or 1 the coefficients of the PCA depending on if they are below or above the average PCA coefficient values.

Moreover, this paper performs an $\epsilon$-clustering (complexity $\mathcal{O}(n^2)$!) to find clusters by comparing the distances of the PCA signatures (before quantization), and then another loop to improve the results by merging clusters where representatives have PCA signatures (after quantization) with a low Hamming distance.

I am really surprised that the complexity $\mathcal{O}(n^2)$ worked for them on 2B images, even if they considered 24-bit hashes.

They also found 1/4 of the images being duplicated (icons, ads, ...).

- [D2LV: A Data-Driven and Local-Verification Approach for Image Copy Detection (2021)](https://arxiv.org/pdf/2111.07090.pdf):
The authors are the winners of the Image Similarity Challenge proposed by Facebook. They used neural networks with an unsupervised pretraining, followed by a training with both a triplet and a cross entropy loss. They used a combination of global and local features (it is often the case in other approaches).

- [3rd Place: A Global and Local Dual Retrieval Solution to Facebook AI Image Similarity Challenge (2021)](https://arxiv.org/pdf/2112.02373.pdf):
This paper presents a method for checking if two images are similar. It uses a combination of global and local features. The global features are obtained with a Transformer pre-trained on ImageNet and trained with a triplet loss, and the local features are from the SIFT algorithm.

- [Detecting Near-Duplicates for Web Crawling (2007)](https://www2007.org/papers/paper215.pdf):
This paper verifies that the SimHash algorithm (for text) creates hashes such that near duplicates have close Hamming distances. Also, and this is what is most important for us, it gives an efficient solution for the Hamming distance problem to efficiently determine closest simhashes (i.e not in quadratic time) via random projections or permutations + interpolation search.

- [An image signature for any kind of image (2002)](https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.104.2585&rep=rep1&type=pdf):
This presents the Image signature algorithm. We start by cropping the image to essentially remove the constants areas around. Then, we divide the image into a grid of patches. For each patch, we compute the average gray level, and compare it to the neighboring patches. We replace the patch by 8 values, all between -2, -1, 0, 1 and 2, depending on the difference between the level of the patch and its 8 neighbors. We then concatenate all the outputs to obtain a 648-bit array. The difference between two hashes is computed with the euclidean distance.

This algorithm is interesting and seems fast. I like the idea of working with patches. However, the vectors are typically longer (size 648, 10x the usual size of 64).

- [A robust image hashing based on discrete wavelet transform (2017)](https://sci-hub.hkvisa.net/10.1109/icsipa.2017.8120651):
This paper introduces a new hashing method, based on image normalization, DWT and SVD. Some ideas can be taken to improve the simple PHash algorithm.

## Less interesting papers

- [Fast and accurate near-duplicate image elimination for visual sensor networks (2016)](https://journals.sagepub.com/doi/pdf/10.1177/1550147717694172):
Interesting paper, but I find it really similar to [Duplicate Discovery on 2 Billion Internet Images (2013)](https://people.csail.mit.edu/celiu/pdfs/CVPR13-bigvis-dupdiscovery.pdf) in the sense that they are also doing a two-step method with global and then local features. The clustering and nearest neighbors search is also similar. They changed the hash function, and added a PageRank algorithm to find the most relevant image to keep from a cluster once it is formed, but I don't think it really matters.

They provide good metrics for the evaluation.

- [Benchmarking unsupervised near-duplicate image detection (2019)](https://arxiv.org/pdf/1907.02821.pdf):
This paper makes a benchmark of existing methods for image deduplication. It is interesting to understand how to perform an evaluation.

- [Large Scale Image Deduplication (2011)](http://vision.stanford.edu/teaching/cs231a_autumn1213_internal/project/final/writeup/nondistributable/Wen_Paper.pdf):
It is based on PCA to compute the image hash. The PCA is done on a sufficiently large image collection but I am not sure performing a PCA is better than PHash.

- [Secure image deduplication through image compression (2015)](https://sci-hub.hkvisa.net/10.1016/j.jisa.2015.11.003):
It is based on a wavelet-based image compression algorithm called SPIHT. It creates the signature of an image by identifying the significant regions on the image. Interesting idea but I am not sure how this can be better than PHash.

- [Image Deduplication Based on Hashing and Clustering in Cloud Storage (2021)](https://koreascience.kr/article/JAKO202120941694290.pdf):
It presents a hashing function based on DCT (that I don't think it's better than DCT) and does the clustering with K-means, but I don't like this clustering strategy as it is really challenging to find $k$ for images.

- [CE-Dedup: Cost-Effective Convolutional Neural Nets Training based on Image Deduplication (2021)](https://arxiv.org/pdf/2109.00899.pdf):
This paper uses techniques like PHash, DHash, AHash or WHash to first deduplicate images in a training set, and then train a neural network on this to see that performances can be really close when training on the full dataset while reducing the size of the dataset by a large amount. However, it does not use a neural network to do the deduplication.

- [Efficient Cropping-Resistant Robust Image Hashing (2014)](https://sci-hub.hkvisa.net/10.1109/ares.2014.85):
This presents the Crop-resistant hashing algorithm.

It is an old-school method and I’m not convinced being more robust against cropping doesn’t hurt the overall performances.

- [A lightweight virtual machine image deduplication backup approach in cloud environment (2014)](https://sci-hub.hkvisa.net/10.1109/compsac.2014.73):
It is based on the K-means algorithm but I don’t like the approach since we don’t know how to choose $k$ and we have a constraint to fit in the RAM.

- [Clustering-based acceleration for virtual machine image deduplication in the cloud environment (2016)](http://lcs.ios.ac.cn/~zhangzy/paper/JSS2016Xu.pdf):
It is essentially the same paper (also same authors) as A lightweight virtual machine image deduplication backup approach in cloud environment.

- [A duplicate image deduplication approach via Haar wavelet technology (2012)](https://sci-hub.hkvisa.net/10.1109/ccis.2012.6664249):
It is based on the wavelet decomposition instead of the DCT, which seems to perform worse.

- [A High-precision Duplicate Image Deduplication Approach (2013)](http://www.jcomputers.us/vol8/jcp0811-06.pdf):
It is based on the wavelet decomposition instead of the DCT, which seems to perform worse.

# Blog posts

https://content-blockchain.org/research/testing-different-image-hash-functions/

https://www.hackerfactor.com/blog/index.php?/archives/432-Looks-Like-It.html

https://www.hackerfactor.com/blog/index.php?/archives/529-Kind-of-Like-That.html

https://fullstackml.com/wavelet-image-hash-in-python-3504fdd282b5

https://towardsdatascience.com/understanding-locality-sensitive-hashing-49f6d1f6134

https://santhoshhari.github.io/Locality-Sensitive-Hashing/

https://en.wikipedia.org/wiki/Locality-sensitive_hashing

https://mesuvash.github.io/blog/2019/Hashing-for-similarity/

https://keras.io/examples/vision/near_dup_search/

https://drivendata.co/blog/image-similarity-winners

https://www.linkedin.com/pulse/detection-duplicate-images-using-deep-learning-aditya-sharma/

https://medium.com/mlearning-ai/a-scalable-solution-to-detect-duplicate-images-97d431c2726d

https://keras.io/examples/vision/siamese_network/


# Clip distributions - descriptive stat

- SBU Captions
```python
DescribeResult(nobs=10000, minmax=(0.11153904348611832, 0.44991129636764526), mean=0.2874957061290741, variance=0.0016425453395696478, skewness=-0.22512623318313724, kurtosis=0.1512977180455395)
```

- Red Caps
```python
DescribeResult(nobs=10000, minmax=(0.08980361372232437, 0.4210364818572998), mean=0.3082767878524959, variance=0.001230211924011678, skewness=-0.5157219676083339, kurtosis=0.6965278169334876)
```

- LAION 400M
```python
DescribeResult(nobs=10000, minmax=(0.16056129336357117, 0.4760231077671051), mean=0.333618477447331, variance=0.0008586748609226699, skewness=0.7131919650316029, kurtosis=1.668628208211425)
```


# Creating image PMD - points of entry

Some subsets are handled on JZ, and some others are handled on thomas-m4-pmd GCP VM.

To launch the ones on JZ (Conceptual captions and Google WIT), from this folder, launch the slurm job:
```batch
sbatch jz_image_pmd.slurm --mail-type=ALL --mail-user=victor@huggingface.co # The two last arguments are optional
```
The job will expect you to have the `m4` repo under `$WORK/code`, and will save data under `$cnw_ALL_CCFRSCRATCH/general_pmd/image`. You can then upload things to the bucket to save them:
```bash
gsutil -m rsync -r $cnw_ALL_CCFRSCRATCH/general_pmd/image/ gs://science-m4/general_pmd/image/
```

To launch the ones on `thomas-m4-pmd`, from this folder, run the following comamnds:
```bash
mkdir -p $HOME/general_pmd/image
# Optionally, activate your favorite conda env
python pmd.py
gsutil -m rsync -r $HOME/general_pmd/image/ gs://science-m4/general_pmd/image/
```

Once the creation are done, you can sanity check how many images are missing using the script `check_none_ims.py`.

A lot of the subsets require manually downloading files and putting them in the right folder. Note that they are files that are automatically downloaded from `facebook/pmd` (https://huggingface.co/datasets/facebook/pmd), please make sure you have filled the authorization wall so that the download can automatically happen.

|Subset|File location|Where to put and what to do|
|--|--|--|
|LocalizedNarrativesFlickr30K|http://shannon.cs.illinois.edu/DenotationGraph/data/index.html|Download "Flickr 30k images" and decompress the tar.gz into `~/.cache/m4/flickr30k`|

## Tarring the `downloaded_images` folders in `~/.cache/m4/`

```bash
find . -type d -maxdepth 1 -mindepth 1 -exec basename \{} ./ \; | parallel --verbose -j 16 --progress "tar -zcf {1}.tar.gz {1}"
```

## Helper scripts

If you want to know how many images were downloaded in a subfolder:
```bash
find {dataset_name} -type f -regextype egrep -regex "{dataset_name}/downloaded_images/[0-9a-f]{3}/[0-9a-f]{3}/[0-9a-f]{64}" | wc -l
```

If you want to remove all .lock files from a subfolder:
```bash
find {dataset_name} -type f -regextype egrep -regex "{dataset_name}/downloaded_images/[0-9a-f]{3}/[0-9a-f]{3}/[0-9a-f]{64}\.lock" | xargs -I {} rm {}
```

If you want to remove all tmp_files from a subfolder:
```bash
find {dataset_name} -type f -regextype egrep -regex "{dataset_name}/downloaded_images/temp-.*" | xargs -I {} rm {}
```

If you want to tar and split the subfolder (typically before pushing to a bucket):
```bash
tar -cvf - {dataset_name} | split -b 1G -d -a 7 --additional-suffix=.tar - "{dataset_name}_part-"
```

Note: on MAC, `split` has to be replaced with `gsplit` (pleaase install it via `brew install coreutils`)


## Locally
Run the `run_document_ngrams_extraction.sh` script.

## On JZ
On JZ:
- Add to your `~/.bashrc` the following line (custom installation of `jq` and `parallel`):
```bash
export PATH=$PATH:/gpfswork/rech/six/commun/lib/jq-1.5/bin/:/gpfswork/rech/six/commun/lib/parallel/bin/
```

Then, run the slurm script (`sbatch pipe.slurm`).


# Code of Conduct

## Our Pledge

In the interest of fostering an open and welcoming environment, we as
contributors and maintainers pledge to make participation in our project and
our community a harassment-free experience for everyone, regardless of age, body
size, disability, ethnicity, sex characteristics, gender identity and expression,
level of experience, education, socio-economic status, nationality, personal
appearance, race, religion, or sexual identity and orientation.

## Our Standards

Examples of behavior that contributes to creating a positive environment
include:

* Using welcoming and inclusive language
* Being respectful of differing viewpoints and experiences
* Gracefully accepting constructive criticism
* Focusing on what is best for the community
* Showing empathy towards other community members

Examples of unacceptable behavior by participants include:

* The use of sexualized language or imagery and unwelcome sexual attention or
advances
* Trolling, insulting/derogatory comments, and personal or political attacks
* Public or private harassment
* Publishing others' private information, such as a physical or electronic
address, without explicit permission
* Other conduct which could reasonably be considered inappropriate in a
professional setting

## Our Responsibilities

Project maintainers are responsible for clarifying the standards of acceptable
behavior and are expected to take appropriate and fair corrective action in
response to any instances of unacceptable behavior.

Project maintainers have the right and responsibility to remove, edit, or
reject comments, commits, code, wiki edits, issues, and other contributions
that are not aligned to this Code of Conduct, or to ban temporarily or
permanently any contributor for other behaviors that they deem inappropriate,
threatening, offensive, or harmful.

## Scope

This Code of Conduct applies within all project spaces, and it also applies when
an individual is representing the project or its community in public spaces.
Examples of representing a project or community include using an official
project e-mail address, posting via an official social media account, or acting
as an appointed representative at an online or offline event. Representation of
a project may be further defined and clarified by project maintainers.

This Code of Conduct also applies outside the project spaces when there is a
reasonable belief that an individual's behavior may have a negative impact on
the project or its community.

## Enforcement

Instances of abusive, harassing, or otherwise unacceptable behavior may be
reported by contacting the project team at <opensource-conduct@meta.com>. All
complaints will be reviewed and investigated and will result in a response that
is deemed necessary and appropriate to the circumstances. The project team is
obligated to maintain confidentiality with regard to the reporter of an incident.
Further details of specific enforcement policies may be posted separately.

Project maintainers who do not follow or enforce the Code of Conduct in good
faith may face temporary or permanent repercussions as determined by other
members of the project's leadership.

## Attribution

This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html

[homepage]: https://www.contributor-covenant.org

For answers to common questions about this code of conduct, see
https://www.contributor-covenant.org/faq

CODE_OF_CONDUCT.md

Large Concept Models: Language modeling in a sentence representation space


# Contributing to large_concept_model
We want to make contributing to this project as easy and transparent as
possible.

## Pull Requests
We actively welcome your pull requests.

1. Fork the repo and create your branch from `main`.
2. If you've added code that should be tested, add tests.
3. If you've changed APIs, update the documentation.
4. Ensure the test suite passes.
5. Make sure your code lints.
6. If you haven't already, complete the Contributor License Agreement ("CLA").

## Contributor License Agreement ("CLA")
In order to accept your pull request, we need you to submit a CLA. You only need
to do this once to work on any of Facebook's open source projects.

Complete your CLA here: <https://code.facebook.com/cla>

## Issues
We use GitHub issues to track public bugs. Please ensure your description is
clear and has sufficient instructions to be able to reproduce the issue.

Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe
disclosure of security bugs. In those cases, please go through the process
outlined on that page and do not file a public issue.

## License
By contributing to large_concept_model, you agree that your contributions will be licensed
under the LICENSE file in the root directory of this source tree.

CONTRIBUTING.md


# Large Concept Models
## Language Modeling in a Sentence Representation Space

[[Blog]](https://ai.meta.com/blog/meta-fair-updates-agents-robustness-safety-architecture/) [[Paper]](https://ai.meta.com/research/publications/large-concept-models-language-modeling-in-a-sentence-representation-space/)

This repository provides the official implementations and experiments for [Large Concept Models](https://ai.meta.com/research/publications/large-concept-models-language-modeling-in-a-sentence-representation-space/) (**LCM**).

The LCM operates on an explicit higher-level semantic representation,
which we name a "concept". Concepts are language- and modality-agnostic and represent a higher
level idea. In this work, a concept corresponds to a sentence, and we use the [SONAR](https://github.com/facebookresearch/SONAR)
embedding space, which supports up to 200 languages in text and 57 languages in speech. See the list of supported languages [here](https://github.com/facebookresearch/SONAR?tab=readme-ov-file#supported-languages-and-download-links).

## Approach

The LCM is a sequence-to-sequence model in the concepts space trained to perform auto-regressive sentence prediction.
We explore multiple approaches:
- MSE regression (`base_lcm` in this code).
- Variants of diffusion-based generation (we include `two_tower_diffusion_lcm` in this release).
- Models operating in a quantized SONAR space (coming soon).

These explorations are performed using 1.6B parameter models and training data in the order of 1.3T tokens. We include in this repository recipes to reproduce the training and finetuning of 1.6B MSE LCM and Two-tower diffusion LCM. See instructions [below](#usage).

## Installing

### Using UV

The LCM repository relies on fairseq2. If you have `uv` installed on your system, you can install a virtual environment with all the necessary packages by running the following commands:
```bash
uv sync --extra cpu --extra eval --extra data
```

You can also use `uv run` to run the demo commands with the correct environment.

Note that we only provide requirements for `cpu` dependencies, if you want to use GPU support, you will have to choose the variants of torch and fairseq2 that work for your system.
For example for torch 2.5.1 with cuda 1.21, You would do something like:
```
uv pip install torch==2.5.1 --extra-index-url https://download.pytorch.org/whl/cu121 --upgrade
uv pip install fairseq2==v0.3.0rc1 --pre --extra-index-url  https://fair.pkg.atmeta.com/fairseq2/whl/rc/pt2.5.1/cu121 --upgrade
```

Check [fairseq2 variants](https://github.com/facebookresearch/fairseq2?tab=readme-ov-file#variants) for possible variants. Note that LCM currently relies on the release candidate for fairseq2 0.3.0 rc1.

### Using pip

To install with pip, the commands are very similar, but you will have to manage your own environment and make sure to install fairseq2 manually first. For instance, for a `cpu` install.

```bash
pip install --upgrade pip
pip install fairseq2==v0.3.0rc1 --pre --extra-index-url  https://fair.pkg.atmeta.com/fairseq2/whl/rc/pt2.5.1/cpu
pip install -e ".[data,eval]"
```

If [fairseq2](https://github.com/facebookresearch/fairseq2) does not provide a build for your machine, check the readme of that project to build it locally.

## Usage

> [!NOTE]
> If using `uv` prefix all commands with `uv run` to use the environment created by default in `.venv`, e.g.,
> `uv run torchrun --standalone`.
> Alternatively, you can activate the environment once and for all with `source .venv/bin/activate`.

### Preparing data

The LCM can be trained and evaluated using textual data split in sentences and embedded with [SONAR](https://github.com/facebookresearch/SONAR/). We provide a sample processing pipeline that can be used to prepare such training data, you can run it with:

```
 uv run --extra data scripts/prepare_wikipedia.py /output/dir/for/the/data
 ```

 This pipeline shows how to get a dataset from huggingface and process it with SONAR and [SaT](https://arxiv.org/abs/2406.16678). Check out the file for more details on processing your own data. While the script provides an example pulling data from huggingface, we also provide [APIs](https://github.com/facebookresearch/stopes/tree/main/stopes/utils/sharding) to process jsonl, parquet and CSV.

### Datacards

The trainer described below relies on datacards configuring the datasets. These datacards are yaml files with pointers to the dataset files (locally or on s3) and information on its schema. We provide some sample datacards in [`lcm/datacards/datacards.yaml`](https://github.com/facebookresearch/large_concept_model/blob/main/lcm/datacards/datacards.yaml). Once you have processed some data, you can update the datacards with your paths.

#### Fitting a normalizer
To fit a new embedding space normalizer on a given weighted mixture of datasets
one can use the following command :
```bash
python scripts/fit_embedding_normalizer.py --ds dataset1:4 dataset2:1 dataset3:10 --save_path "path/to/new/normalizer.pt" --max_nb_samples 1000000
```
Here, `dataset1`, `dataset2`, `dataset3` are the names of datasets declared in the datacards as shown above
and `(4, 1, 10)` their respective relative weights.
The resulting normalizer can be next declared as a model as shown in `lcm/cards/sonar_normalizer.yaml`
and referenced in all model training configs.

### Pre-training models

#### Base MSE LCM

To train an MSE LCM, we will use one of the following commands:

**Option 1.** Training with SLURM using [submitit](https://github.com/facebookincubator/submitit) via [stopes](https://github.com/facebookresearch/stopes/tree/main)'s launcher:
```sh
python -m lcm.train \
    +pretrain=mse \
    ++trainer.output_dir="checkpoints/mse_lcm" \
    ++trainer.experiment_name=training_mse_lcm \
```
With this command, we will submit a slurm job named `training_mse_lcm` with the recipe's requirements, in this case:
```yaml
requirements:
  nodes: 4
  tasks_per_node: 8
  gpus_per_node: 8
  cpus_per_task: 32
  mem_gb: 0
  timeout_min: 10000
```
You can override the job's requirements like the timeout limit and the launcher's slurm partition with:
```sh
python -m lcm.train \
    +pretrain=mse \
    ++trainer.output_dir="checkpoints/mse_lcm" \
    ++trainer.experiment_name=training_mse_lcm \
    ++trainer.requirements.timeout_min=100 \
    ++trainer.requirements.cpus_per_task=8 \
    ++launcher.partition=$partition_name
```

**Option 2.** Training locally with `torchrun` (e.g. using only 2 GPUs) with a smaller batch size (overriding `++trainer.data_loading_config.max_tokens=1000`):
```sh
CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nnodes=1 --nproc-per-node=2 \
    -m lcm.train launcher=standalone \
    +pretrain=mse \
    ++trainer.data_loading_config.max_tokens=1000 \
    ++trainer.output_dir="checkpoints/mse_lcm" \
    +trainer.use_submitit=false \
```
> [!IMPORTANT]
> Since we're changing the number of GPUs required by the recipe, this will not reproduce the experimental setup of the paper.

The checkpoints directory `checkpoints/mse_lcm` will be structured as:
```
.
├── checkpoints
│   ├── step_2000
│   ├── ...
│   └── step_250000
├── config_logs
├── executor_logs
├── model_card.yaml
├── tb   # tensorboard logs
└── wandb  # W&B logs
```
Note that W&B logging is skipped unless `wandb` is available.
You can install `wandb` with `uv pip install wandb`.
W&B arguments can be changed by overriding Hydra config values in the recipe:

```sh
++trainer.wandb_project=$project_name
++trainer.wandb_run_name=$run_name
```

#### Two-tower diffusion LCM

Similar to the base MSE LCM we can submit a training job following the recipe in [./recipes/train/pretrain/two_tower.yaml](./recipes/train/pretrain/two_tower.yaml) via:

```sh
python -m lcm.train \
    +pretrain=two_tower \
    ++trainer.output_dir="checkpoints/two_tower_lcm" \
    ++trainer.experiment_name=training_two_tower_lcm \
```

> [!TIP]
> To understand the different ingredients of training recipes, check [this README](./recipes/train/README.md).

### Finetuning models
To finetune the previously pre-trained two-tower diffusion LCM on supervised data,  follow these steps:

**Step 1.** Register the pre-trained checkpoint as a fairseq2 asset.

You can finetune the final checkpoint with the card `checkpoints/two_tower_lcm/model_card.yaml` or any checkpoint after a specific number of training steps, e.g., `checkpoints/two_tower_lcm/checkpoints/step_2000/model_card.yaml`.
To register the selected checkpoint, copy the automatically created yaml file to `./lcm/cards/mycards.yaml` and rename the model to replace the default `on_the_fly_lcm`.
`./lcm/cards/mycards.yaml` will look like:
```yaml
__source__: inproc
 checkpoint: file://path_to/large_concept_model/checkpoints/two_tower_lcm/checkpoints/step_2000/model.pt
 model_arch: two_tower_diffusion_lcm_1_6B
 model_family: two_tower_diffusion_lcm
 name: my_pretrained_two_tower
```
For more on how to manage fairseq2 assets, see [documentation](https://facebookresearch.github.io/fairseq2/nightly/basics/assets.html).

**Step 2.** Launch a finetuning job pointing to the model to finetune, in this instance `my_pretrained_two_tower`:
```sh
CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nnodes=1 --nproc-per-node=2 \
    -m lcm.train launcher=standalone \
    +finetune=two_tower \
    ++trainer.output_dir="checkpoints/finetune_two_tower_lcm" \
    ++trainer.data_loading_config.max_tokens=1000 \
    +trainer.use_submitit=false \
    ++trainer.model_config_or_name=my_pretrained_two_tower
```
or

```sh
python -m lcm.train \
    +finetune=two_tower \
    ++trainer.output_dir="checkpoints/finetune_two_tower_lcm" \
    ++trainer.experiment_name=finetune_two_tower_lcm \
    ++trainer.model_config_or_name=my_pretrained_two_tower
```

Similarly, to finetune an MSE LCM, follow the same instructions for registering a pre-trained checkpoint and submit a finetuning job with the appropriate recipe ([./recipes/train/finetune/mse.yaml](./recipes/train/finetune/mse.yaml)) via:
```sh
python -m lcm.train \
    +finetune=mse \
    ++trainer.output_dir="checkpoints/finetune_mse_lcm" \
    ++trainer.experiment_name=finetune_mse_lcm \
    ++trainer.model_config_or_name=my_pretrained_mse_lcm
```
### Evaluating models

> [!NOTE]
> For advanced evaluation (benchmarking different tasks, comparing results with LLMs, etc.) , check [the evaluation documentation](./examples/evaluation/README.md).

**Step 0.** Download NLTK data required for evaluating ROUGE:
```py
python -m nltk.downloader punkt_tab
```

**Step 1.**
Generate and score outputs of a model either by pointing to its `model_card` yaml file or after registering it as a fairseq2 asset (the same way we registerd `my_pretrained_two_tower`):
```sh
model_card=./checkpoints/finetune_two_tower_lcm/checkpoints/step_1000/model_card.yaml
OUTPUT_DIR=evaluation_outputs/two_tower

torchrun --standalone --nnodes=1 --nproc-per-node=1 -m lcm.evaluation  \
  --predictor two_tower_diffusion_lcm  \
  --show_progress true \
  --data_loading.max_samples 100 \
  --model_card ${model_card} \
  --launcher standalone \
  --dataset.source_suffix_text '[MODEL]:' \
  --tasks finetuning_data_lcm.validation \
   --task_args '{"max_gen_len": 10, "eos_config": {"text": "End of text."}}' \
  --data_loading.batch_size 4  --generator_batch_size 4 \
  --dump_dir ${OUTPUT_DIR} \
  --inference_timesteps 40 \
  --initial_noise_scale 0.6 \
  --guidance_scale 3 \
  --guidance_rescale 0.7
```
where in the example we are evaluating 100 samples only (`--data_loading.max_samples 100`) and limiting the model output length to 10 sentences (`--task_args '{"max_gen_len": 10}'`).

Outputs dumped in `./evaluation_outputs/two_tower` will be structured as:
```
.
├── metadata.jsonl
├── metrics.eval.jsonl
├── raw_results
├── results
└── tb
```
where `metrics.eval.jsonl` contains corpus-level scores.

To evaluate an MSE LCM, we use the associated predictor (`base_lcm`) and evaluate with:

```sh
model_card=./checkpoints/finetune_mse_lcm/checkpoints/step_1000/model_card.yaml
OUTPUT_DIR=evaluation_outputs/mse_lcm

torchrun --standalone --nnodes=1 --nproc-per-node=1 -m lcm.evaluation  \
  --predictor base_lcm --sample_latent_variable False \
  --show_progress true \
  --data_loading.max_samples 100 \
  --model_card ${model_card} \
  --launcher standalone \
  --dataset.source_suffix_text '[MODEL]:' \
  --tasks finetuning_data_lcm.validation \
   --task_args '{"max_gen_len": 10, "eos_config": {"text": "End of text."}}' \
  --data_loading.batch_size 4  --generator_batch_size 4 \
  --dump_dir ${OUTPUT_DIR} \
```

Note that in this example, we only show how to evaluate the LCM on the same finetuning dataset (validation split). To evaluate in a downstream task, and compare results with the LLM, refer to the [Evaluation documentation](./examples/evaluation/README.md).

## Contributing

See the [CONTRIBUTING](CONTRIBUTING.md) file for how to help out.

## Citation

If you use this codebase, please cite:
```
@article{lcm2024,
  author = {{LCM team}, Lo\"{i}c Barrault, Paul-Ambroise Duquenne, Maha Elbayad, Artyom Kozhevnikov, Belen Alastruey, Pierre Andrews, Mariano Coria, Guillaume Couairon, Marta R. Costa-juss\`{a}, David Dale, Hady Elsahar, Kevin Heffernan, Jo\~{a}o Maria Janeiro, Tuan Tran, Christophe Ropers, Eduardo Sánchez, Robin San Roman, Alexandre Mourachko, Safiyyah Saleem, Holger Schwenk},
  title = {{Large Concept Models}: Language Modeling in a Sentence Representation Space},
  publisher = {arXiv},
  year = {2024},
  url = {https://arxiv.org/abs/2412.08821},
}
```

## License

This code is released under the MIT license (see [LICENSE](./LICENSE)).

README.md


## Why ?

Why do we need to implement this feature ? What is the use case ?

## How ?

Document the technical decisions you made.
If some parts are WIP, please explicit them here.

## Test plan

How did you test your changes ?
Include full command line to help other people reproduce if needed.

.github/pull_request_template.md


# Evaluation

After you have trained an LCM, the checkpoint will be saved in a folder under the name `model.pt`, together with the model card under the name `model_card.yaml`. We also provide the library to evaluate the LCM and LLM. Using this library brings many benefits: You can reproduce the experiments done in the paper, you can inspect the results in an unified way, and you can also scale up the experiments for very large datasets in SLURM cluster. This document shows how to evaluate the model for different downstream tasks using the LCM eval library.

## Step 1: Prepare the data

Since an LCM expects input data in sentence level, we need to preprocess the evaluation datasets accordingly. This includes parsing the raw content and
splitting  texts into sentences, then embedding them into vectors using a Sonar encoder.

The example below shows how we prepare the data for CNN Dailymail. We load the dataset from Huggingface using [`datasets` API](https://huggingface.co/docs/datasets/en/index). The sentence splitting is done using [wtpsplit](https://github.com/segment-any-text/wtpsplit). Make sure to specify `--extra data` in installing the project to include these libraries.

All processing logic is implemented in the file `prepare_evaluation_data.py`, as described below.

### Step 1.1: Process the split:
Next, we download and parse the content (source text and summaries), saving different splits into JSON format

```shell
uv run --extra data prepare_evaluation_data.py prepare_data \
    --dataset_name=cnn_dailymail \
    --output_dir=jsonl_dataset \
    --source_text_column=article \
    --target_text_column=highlights \
    --version=3.0.0 \
    --prompt_prefix="Summarize the following news to a concise list of highlights.\n[Text Start]:\n"
    --prompt_suffix="\n[Text End]"
```

Explain: In the above script, `cnn_dailymail` and `3.0.0` is the name and configuration of the dataset as available in HuggingFace `datasets`, `article` and `highlights` are source and summary columns. The `prompt_prefix` and `prompt_suffix` are optional arguments, if specified they will be prepended and appended to each source text to form the complete prompt. These arguments are useful if you want to embed the prompts into the dataset, and let them process all at once together with the text. Alternatively, we can specify them at later phase, when we evaluate the model (in which case the model will process the prompts on the fly)

> **_NOTE:_**  When `prompt_prefix` or `prompt_suffix` are specified, the dataset schema will change, i.e. the columns are renamed to "prompt" for input and "answer" for output. This is to indicate that we are handling the "processed" dataset and not the original one.

The output will be stored in different files `[split].jsonl` under the directory `output_dir`.

### Step 1.2: Sentence splitting and embedding:

To perform sentence splitting and sonar embedding for each split, run the following command:

```shell
uv run --extra data prepare_evaluation_data.py embed \
    --input_path=jsonl_dataset/cnn_dailymail/test.jsonl \
    --source_text_column=prompt \
    --target_text_column=answer \
    --output_dir=parquet_dataset/cnn_dailymail \
    --lang=eng_Latn \
    --mode=local \
    --log_dir=/tmp/logs/embed_cnndm
```

Depending on your machine, this might take some time. Alternatively, you can try to run in your SLURM cluster with the argmnent `--mode=slurm --shards=NO_OF_PARALLEL_JOBS`. This requires changing your SLURM config accordingly. We use [submitit](https://github.com/facebookincubator/submitit) to configure the job launcher. Here is the relevant excerpt in the script:

```python
launcher = Launcher(
      cache=None,
      config_dump_dir=Path(log_dir) / "conf",
      log_folder=Path(log_dir) / "logs",
      cluster=mode,
      update_parameters={"partition": "your_slurm_partition"},
  )
  _ = await launcher.schedule(inst_stopes_module)
```

## Step 2: Choose the predictor for evaluation

To run the evaluation, we first need to map the model to a `Predictor`, which is an object that streamlines a number of steps: Loading the models, reading the prompts, performing the inference, decoding the outputs according to a given user setting, and finally formatting the text into the user-friendly format. Currently, the list of supported model families and their predictors is below. All predictors are found in "lcm/evaluation/predictors" and are registered in lcm.evaluation.predictors`_PREDICTOR_CONFIG_MAP`

| Predictor               | Model family           | Model identifier                                           |
|-------------------------|------------------------|------------------------------------------------------------|
| huggingface             | AutoModel transformers | `model_name`, `revision`, `model_class`, `tokenizer_class` |
| llama3                  | LlaMA 3.x              | `model_name`                                               |
| gemma                   | Gemma                  | `model_name`                                               |
| base_lcm                | Base LCM               | `model_card`                                               |
| two_tower_diffusion_lcm | Two-tower diffusion LCM| `model_card`                                               |

Next, we specify how the decoder generate texts with different generation options.

For LLMs, the options are parameters found in [transformers.GenerationConfig](https://huggingface.co/docs/transformers/en/main_classes/text_generation#transformers.GenerationConfig), and we port in the predictors the most popular ones: `repetition_penalty`, `encoder_repetition_penalty`, `encoder_no_repeat_ngram_size`, `no_repeat_ngram_size`.

For LCMs, the options are found in [LCMGeneratorOptions](https://github.com/facebookresearch/large_concept_model/blob/main/lcm/inference/lcm/generator.py#L31) (for Base LCM) or [DiffusionLCMGeneratorOptions](https://github.com/facebookresearch/large_concept_model/blob/main/lcm/inference/two_tower_diffusion_lcm/generator.py#L31) (for Two-tower diffusion LCM). These options only specify how to generate output embeddings using diffusion process. We also want to specify the sonar decoder options, which dictates how the embeddings are decoded into texts, using parameters in [SonarDecoderConfig](https://github.com/facebookresearch/large_concept_model/blob/main/lcm/datasets/configs.py#L69).

## Step 3: Choose a downstream task and run the evaluation

To run the downstream task, specify the task name and configuration, as well as parameters. We provide example tasks that were used in the paper:

### LLM evaluation tasks:

| Task name               | Task configuration               | Explanation                                                                                |
|-------------------------|----------------------------------|-------------------------------------------------------------------------------------------------------------|
| cnn_dailymail           | cnn_dailymail_{form}llm.{split}  | {form} can be empty for or "inverse_" for summary expansion, {split} can be "test", "validation" or "train" |
| xsum                    | xsum_{form}llm.{split}           | {form} can be empty for or "inverse_" for summary expansion, {split} can be "test", "validation" or "train" |
| xlsum_llm               | xlsum_llm.{lang}.{split}         | {lang} refers to one value in [language list](../../lcm/evaluation/tasks/xlsum.py), {split} can be "test", "validation" or "train" |

The evaluation library provides the handy CLI to evaluate using `lcm.evaluation` entry. Example command for evaluating the Meta Llama 3.1 8B instruction:

```shell
uv run torchrun --standalone --nnodes=1 --nproc-per-node=1 -m lcm.evaluation \
  --predictor llama3  \
  --model_name meta-llama/Llama-3.1-8B-Instruct \
  --generator_batch_size 16 \
  --tasks cnn_dailymail_llm.test \
  --task_args '{"max_gen_len": 200}' \
  --dataset_dir jsonl_dataset/cnn_dailymail \
  --data_loading.batch_size 16 \
  --dataset.soure_text_column prompt \
  --dataset.source_target_column answer \
  --dump_dir output_results
```

In the example above, we load the model "meta-llama/Llama-3.1-8B-Instruct" as [specified](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) in HuggingFace, evaluate it on the CNN dailymail in which we process using the `prepare_evaluation_data.py` script as in Step 1.1, and store the results in the folder specified via `dump_dir`. The argument `dataset_dir` refers to the value of the argument `output_dir` in Step 1.1.

In some cases, the model requires authentication token to evaluate. You can obtain them in HuggingGface (see [User Access Tokens](https://huggingface.co/docs/hub/en/security-tokens)), then add the parameter `--use_auth_token [YOUR TOKEN]` to the CLI command.

In the above example, we need to provide the `source_text_column` and `source_target_column` parameters, because in Step 1, we inject the prompts direcly to the dataset and renamed the columns accordingly (to differentiate with "original" datasets).  You can also skip this part and customize the prompt for each for each evaluation run. To do this, instead of specifying the `prompt_prefix` and `prompt_suffix` when preparing the data (as shown in the example in Section 1.1), we specify `dataset.source_prefix_text` and `dataset.source_suffix_text` during the evaluation run:

```shell
uv run torchrun --standalone --nnodes=1 --nproc-per-node=1 -m lcm.evaluation \
  --predictor llama3  \
  --model_name meta-llama/Llama-3.1-8B-Instruct \
  --generator_batch_size 16 \
  --tasks cnn_dailymail_llm.test \
  --task_args '{"max_gen_len": 200}' \
  --dataset_dir jsonl_dataset/cnn_dailymail \
  --data_loading.batch_size 16 \
  --dataset.source_prefix_text "Summarize the following news to a concise list of highlights.\n[Text Start]:\n" \
  --dataset.source_suffix_text "\n[Text End]" \
  --dump_dir output_results
```

> **_NOTE:_** the missing parameters `source_text_column` and `target_text_column` and the new parameters `source_prefix_text`, `target_prefix_text` are becase we do not modify the column schema. Therefore, the original text columns ("article", "highlights") are kept and not specified in the CLI.

It is also possible to provide the prompt from a YAML file. This is handy when you have to engineer the prompts carefully and have a very long detailed text. We provide one example prompt in the file [instruction.yaml](./instruction.yaml). The example command is:

```shell
uv run torchrun --standalone --nnodes=1 --nproc-per-node=1 -m lcm.evaluation \
  --predictor llama3  \
  --model_name meta-llama/Llama-3.1-8B-Instruct \
  --generator_batch_size 16 \
  --tasks cnn_dailymail_llm.test \
  --task_args '{"max_gen_len": 200}' \
  --dataset_dir jsonl_dataset/cnn_dailymail \
  --data_loading.batch_size 16 \
  --prompt_file instruction.yaml \
  --dump_dir output_results
```

### LCM evaluation tasks:

In contrast to LLM, the LCMs expect dataset to be preprocessed in Parquet format, with inputs being (sonar-) sentence embeddings. To evaluate an LCM on a ddownstream task, point to the directory consisting of the parquet files, as specified in Step 1, and run (example for Two-tower diffusion LCM):

```shell
uv run torchrun --standalone --nnodes=1 --nproc-per-node=1 -m lcm.evaluation \
  --predictor two_tower_diffusion_lcm  \
  --model_card path/to/the/model_card.yaml \
  --generator_batch_size 16 \
  --tasks lcm_generation \
  --task_args '{"max_gen_len": 200}' \
  --dataset.parquet_path parquet_dataset/cnn_dailymail \
  --dataset.source_column prompt_sentences_sonar_emb \
  --dataset.source_text_column prompt_sentences \
  --dataset.target_column answer_sentences_sonar_emb \
  --dataset.target_text_column answer_sentences \
  --data_loading.batch_size 16 \
  --dump_dir output_results
```

Similar to LLM evaluation, it is possible to specify the prompt prefix and suffix ad-hoc. This text will be sentence-split and embedded using the standard Sonar encoder.

## Common CLI arguments

|  Argument | Description |
|----------|----------|
| `predictor` | The wrapper of the nodel to be evaluated. See Step 2 for more details
| `data_loading.max_samples`   | Evaluate on the maximum _k_ examples in the test data. Useful for debugging   |
| `data_loading.batch_size`   | Loading and evaluate data in batch. By default `batch_size=10`   |
| `dataset_dir` | The directory consists of different JSONL files processed in Step 1. Only used in LLM evaluation
| `dataset.parquet_path` | The parquet path  consists of different Parquet files files processed in Step 1. Only used in LCM evaluation
| `dataset.source_column` | The column in the data that refers to the input embedding. Not applicable when evaluating LLMs.
| `dataset.source_text_column` | The column in the data that refers to the input text.
| `dataset.target_column` | The column in the data that refers to the ground-truth embedding. Not applicable when evaluating LLMs.
| `dataset.target_text_column` | The column in the data that refers to the ground-truth text.
| `dataset.source_text_prefix` | The text that will prepended to each input text to make the prompt for the model.
| `dataset.source_text_suffix` | The text that will appended after each input text to make the prompt for the model.
| `task_args` | The JSON-formatted string that represents the task arguments. See [task param list](#task_param_list) below.
| `dump_dir` | The directory consisting output of the eval run. If successful, there should be a file `metrics.eval.jsonl` that consists of metric results, the directory `results` that capture the verbose command line used with the detailed output scores, and the directory `raw_results` that shows
the model output for each individual sample, together with the per-sample metric results.
| `task` | Task configuration. See Step 3 for examples.
| `task_args` | The JSON-formatted string that represents the task arguments. See [task param list](#task_param_list) below.
| `launcher` | Whether the CLI should be run locally, or in SLURM cluster. Accepted value is `local`, `submitit` (SLURM) or `standalone` (debug mode).
| `job_args` | Parameters used when launching eval in SLURM. See [below](#slurm-eval) for more details.

*Table: List of common arguments in Evaluation CLI.*

_Note_: In above examples, free arguments such as `generator_batch_size`, `temperature`, etc. are generator options. They depend on specific predictor, as explained in Step 2. Giving a wrong option will trigger and error in the CLI.

Outputs dumped in the directory specified by `dump_dir` will be structured as:
```
.
├── metadata.jsonl
├── metrics.eval.jsonl
├── raw_results
├── results
└── tb
```
where `metrics.eval.jsonl` contains corpus-level scores.

### Task arguments

In both LLM and LCM evaluation, we can configure how inputs and outputs are processed:

- `max_prompt_len`: The model context size, i.e. maximum number of tokens (in LLM) or sentences (in LCM) that the model can accept
- `max_gen_len`: The maximum number of tokens (in LLM) or sentences (in LCM) the model should generate. Note that some model generators have its own stopping criteria, so the actual generated text can be much lower than this value.
- `min_gen_len`: The minimum number of tokens (in LLM) or sentences (in LCM) the model should generate.
- `max_gen_len_ratio`: The maximum  number of tokens (in LLM) or sentences (in LCM) the model should generate _in comparison_ to the input length. For example, if the source document is 5K long and `max_gen_len_ratio=0.2`, we are asking the model to generate 1K-long output (Again, due to the model generators inner behaviour, the output can be much shorter)

## Evaluate big datasets

The above command is sufficient for most cases where you load the model into one GPU and evaluate the whole dataset locally, i.e. the datasets and everyhing is loaded into the memory.
For bigger datasets, or for models which are not easily run in one GPU, or two slow to evaluate, we can submit the evaluation job to the SLURM cluster by choosing the `launcher=submitit`:

```shell

slurm_partition=YOUR_SLURM_PARTITION
shards=NUMBER_OF_SLURM_NODES
timeout_min=JOB_TIMEOUT_IN_MINUTES

uv run -m lcm.evaluation \
  --predictor two_tower_diffusion_lcm  \
  --model_card path/to/the/model_card.yaml \
  --generator_batch_size 16 \
  --tasks lcm_generation \
  --task_args '{"max_gen_len": 200}' \
  --dataset.parquet_path parquet_dataset/cnn_dailymail \
  --data_loading.batch_size 16 \
  --dump_dir output_results \
  --launcher submitit \
  --job_args '{"launcher.cache": "null", "launcher.partition": "'${slurm_partition}'", "launcher.qos": "'${qos}'", "nshards": '${shards}', "requirements": {"gpus_per_node": 1, "timeout_min": '${timeout_min}'}}' \
```

The parameters in `job_args` are submitit parameters. Please refer to https://github.com/facebookincubator/submitit for more comprehensive documentation and parameters list.


# Main ingredients of training recipes

### Training and validation data

```yaml
training_data:
  - name: "<corpus_name>=<split>:<weight>"
    source_prefix_text: "Beginning of source."  # concept added at the beginning of source
    source_suffix_text: "End of source."        # concept added at the end of source
    target_prefix_text: "Beginning of target."  # concept added at the beginning of target (supervised data only)
    target_suffix_text: "End of target."        # concept added at the end of target (supervised data only)

  - name: "<corpus2_name>=<split>:<weight2>"

```

### Data loading config

```yaml
data_loading_config:
  max_tokens: 7168      # Exclusive with batch_size
  batch_size: none        # Exclusive with max_tokens
  len_to_wrap_long_seq: 128 # Sequences longer than this will be wrapped.
  packing: true     # if True, documents in the batch will be packed.
```

The batch content can be defined in several ways:
- `max_tokens` / `len_to_wrap_long_seq` approximate `batch_size`.
- `batch_size` x `len_to_wrap_long_seq` approximate `max_tokens`.

Note that `len_to_wrap_long_seq` has to be smaller than the model's `max_seq_len` defined in the architecture (e.g. [`two_tower_diffusion_lcm_1_6B`](../../lcm/models/two_tower_diffusion_lcm/archs.py#L36)`).

To filter out long samples without wrapping, you can add `filters` to each dataset config to filter based on the length of the document's list of sentences (`text_sentences`):
 ```yaml
  - name: "<corpus_name>=<split>:<weight>"
    source_prefix_text: "Beginning of source."
    filters: 'pa.compute.less(pa.compute.list_value_length(pa.dataset.field("text_sentences")), 128)'
```
### Checkpointing config

```yaml
checkpoint_every_n_steps: 2_000     # QED
keep_last_n_checkpoints: 2          # delete all but last N non-consolidated checkpoints
save_model_every_n_steps: 10_000    # consolidate model every N steps (valid if using FSDP)
preserve_consolidated_models: True  # preserve the consolidated checkpoints
```