chonky_distilbert-base-multilingual-cased

Model Details

Model Description

Fine-tune of distilbert/distilbert-base-multilingual-cased trained on nearly 11 billion tokens from more than 34 million Wikipedia articles to predict paragraph breaks. This model can be used to split arbitrary natural language texts into semantic chunks.

Model Sources

• Repository: https://github.com/mamei16/chonky
• Demo: https://huggingface.co/spaces/mamei16/chonky_chunk

Uses

This model can be used as part of a RAG pipeline to hopefully improve downstream performance.

Bias, Risks, and Limitations

This model has been fine-tuned on non-fictional natural language from Wikipedia. As such, it may not work as well on fictional texts containing dialog or poems, mathematical expressions or code.

How to Get Started with the Model

pip install git+https://github.com/mamei16/chonky

Usage:

from chonky import ParagraphSplitter

splitter = ParagraphSplitter(device="cpu", model_id="mamei16/chonky_distilbert-base-multilingual-cased")

text = """Before college the two main things I worked on, outside of school, were writing and programming. I didn't write essays. I wrote what beginning writers were supposed to write then, and probably still are: short stories. My stories were awful. They had hardly any plot, just characters with strong feelings, which I imagined made them deep. The first programs I tried writing were on the IBM 1401 that our school district used for what was then called "data processing." This was in 9th grade, so I was 13 or 14. The school district's 1401 happened to be in the basement of our junior high school, and my friend Rich Draves and I got permission to use it. It was like a mini Bond villain's lair down there, with all these alien-looking machines — CPU, disk drives, printer, card reader — sitting up on a raised floor under bright fluorescent lights."""

for chunk in splitter(text):
  print(chunk)
  print("--")

# Before college the two main things I worked on, outside of school, were writing and programming. I didn't write essays. I wrote what beginning writers were supposed to write then, and probably still are: short stories. My stories were awful. They had hardly any plot, just characters with strong feelings, which I imagined made them deep.
# --
# The first programs I tried writing were on the IBM 1401 that our school district used for what was then called "data processing." This was in 9th grade, so I was 13 or 14. The school district's 1401 happened to be in the basement of our junior high school, and my friend Rich Draves and I got permission to use it. It was like a mini Bond villain's lair down there, with all these alien-looking machines — CPU, disk drives, printer, card reader — sitting up on a raised floor under bright fluorescent lights.
# --

Training Details

Training Data

Link: https://huggingface.co/datasets/mamei16/multilingual-wikipedia-paragraphs

Note that the data has been pre-tokenized and truncated using the tokenizer from distilbert/distilbert-base-multilingual-cased.

The training data is based on wikimedia/wikipedia. Although it's claimed that unwanted sections such as "References", "See more" etc. have been removed from that dataset, that is not the case. Hence, a semi-accurate and half-manual procedure was used to identify and remove those sections. This procedure involved finding the most common short paragraphs in the last 500 characters of the articles in each language and translating each of them to identify the section headers corresponding to the English "References" etc.

Another issue was the existence of several million articles, most notably in the languages Waray, Cebuano and Swedish, which had been written by a bot named Lsjbot. To not artificially inflate low-resource languages and to ensure that most (if not all) articles were written by humans, the offending articles were removed.

Lastly, articles containing too many extremely long or short paragraphs were removed, as well as stub articles.

Training Procedure

Preprocessing

All articles were truncated to max. 512 tokens.

For designing the training strategy, the following desirable traits were identified:

• Avoid catastrophic forgetting
• Avoid sudden drastic changes in distribution
• Boost low-resource languages
• Uphold good performance on high-resource/most commonly used languages
• Must be implementable with static training data

In the end, the training data was generated by combining the datasets of all 104 languages using temperature sampling without replacement and with τ = 2. This would boost low-resource languages, and languages would "die out" one by one spread over the entire epoch (see figure below), thus avoiding sudden big changes in distribution. Furthermore, the model would keep seeing high-resource languages until the end, thus making it likely that it would maintain good performance on them. At the same time, the linear learning rate schedule would ensure that the more low-resource languages were exchausted during training, the lower the learning rate would be, thus making catastrophic forgetting less likely.

<figure> <img alt="Image" src="assets/languages_remaining.png" /> <figcaption>Fig.1 - Number of languages remaining vs number of training steps taken</figcaption> </figure>

Training Hyperparameters

• Training regime: fp16 mixed precision
• Epochs: 1
• Batch size: 64
• Start learning rate: 2e-5
• Optimizer: Adam
• Weight decay: 0.01
• Loss: NLLLoss
• Label Smoothing factor: 0.1

Evaluation

​MTCB Nano Benchmark (Aggregated Score)

Note: This benchmark is English only (and includes code in multiple programming languages)

Score = Mean of mean_recall, mean_precision, mean_mrr, and mean_ndcg across k=[1, 3, 5, 10] (Metrics reference)

Model Implementation Details

Testing Data

The testing data can be found in the "test" split of each language in the dataset.

Testing Metrics

Due to the extreme class imbalance, the F1-score was chosen as the main evaluation metric.

Results

Mean F1 Score: 0.824

Technical Specifications

Hardware

RTX 5090