vsabolcec's picture
Super-squash branch 'main' using huggingface_hub
53bccb8
|
Raw History Blame Contribute Delete
4.77 kB
metadata
license: apache-2.0
language:
  - eng
  - deu
  - fra
  - pol
  - por
  - spa
  - ita
  - cmn
  - nld
  - afr
  - als
  - amh
  - arb
  - ars
  - ary
  - arz
  - asm
  - azj
  - bel
  - ben
  - bew
  - bod
  - bos
  - bul
  - cat
  - ces
  - ckb
  - cym
  - dan
  - div
  - ekk
  - ell
  - epo
  - eus
  - fas
  - fil
  - fin
  - gle
  - glg
  - gmh
  - guj
  - heb
  - hif
  - hin
  - hrv
  - hun
  - hye
  - ind
  - isl
  - jpn
  - kan
  - kat
  - kaz
  - khk
  - khm
  - kir
  - kmr
  - kor
  - lao
  - lat
  - lit
  - ltz
  - lvs
  - mal
  - mar
  - mkd
  - mlt
  - mya
  - nno
  - nob
  - npi
  - nrm
  - ory
  - pan
  - pbt
  - plt
  - ron
  - rus
  - sin
  - slk
  - slv
  - snd
  - som
  - srp
  - srp
  - swe
  - swh
  - tam
  - tat
  - tel
  - tgk
  - tha
  - tur
  - uig
  - ukr
  - urd
  - uzn
  - uzn
  - vie
  - ydd
  - zsm

FineWeb2-HQ-plus-Classifier

This repository contains the model weights of the trained deep learning classifier used to identify high-quality text samples for the FineWeb2-HQ-plus dataset.

FineWeb2-HQ-plus is a high-quality, model-filtered, multilingual pretraining dataset derived as a subset FineWeb2 and an improvement over FineWeb2-HQ covering 100 languages. The classifier uses XLM-RoBERTa embeddings to score the documents and supports English and additional 100 languages.

For more details, see our paper Toward Cross-Lingual Quality Classifiers for Multilingual Pretraining Data Selection.

Quickstart

Classifier uses a simple architecture that takes mean-pooled XLM-RoBERTa embeddings as input and outputs the score logit.

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer
import huggingface_hub


class BinaryClassifier(torch.nn.Module):
    def __init__(self, embedding_dim=768, hidden_dim=256):
        super(BinaryClassifier, self).__init__()
        self.classifier = torch.nn.Sequential(
            torch.nn.Linear(embedding_dim, hidden_dim),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.2),
            torch.nn.Linear(hidden_dim, 1),
        )

    def forward(self, X):
        return self.classifier(X)

    def to_pt(self, file_name):
        torch.save(self.state_dict(), file_name)

    @classmethod
    def from_pt(cls, file_name, embedding_dim=768, hidden_dim=256):
        state_dict = torch.load(
            file_name,
            weights_only=True,
            map_location=torch.device("cpu"),
        )
        classifier = BinaryClassifier(
            embedding_dim=embedding_dim,
            hidden_dim=hidden_dim,
        )
        classifier.load_state_dict(state_dict)
        classifier.eval()
        return classifier


if __name__ == "__main__":
    embedding_model_name = "FacebookAI/xlm-roberta-base"
    tokenizer = AutoTokenizer.from_pretrained(embedding_model_name)
    embedding_model = AutoModel.from_pretrained(
        embedding_model_name,
        dtype=torch.bfloat16,
    )

    classifiers_dir = huggingface_hub.snapshot_download("epfml/FineWeb2-HQ-plus-Classifier")
    classifier_model = BinaryClassifier.from_pt(f"{classifiers_dir}/model.pt")

    def score_sample(text, classifier_model):
        inputs = tokenizer([text], return_tensors="pt")
        embeddings = embedding_model(**inputs).last_hidden_state.float().mean(1)
        score = F.sigmoid(classifier_model(embeddings)).item()
        return score

    text_en = "Question: How is bipolar disorder different from unipolar depression or 'regular' depression?\nAnswer: Both bipolar disorder and major depression are typically associated with depressive episodes. So both illnesses are accompanied by depressions. The difference is that in bipolar disorder people also have periods of elevation -- or severe irritability. We call these manic or hypomanic episodes."
    score = score_sample(text_en, classifier_model)
    print(f"{score:0.4f}")  # 0.9952

    text_en = "Custom Wedding Gifts\nPersonalized photo frames, albums & keepsakes. Heirloom quality!\nCustom Engraved Journals\nHandmade in Florence Italy. Dozens of sizes and paper styles!"
    score = score_sample(text_en, classifier_model)
    print(f"{score:0.4f}")  # 0.0001

Citation information

@misc{turki2026crosslingualqualityclassifiersmultilingual,
      title={Toward Cross-Lingual Quality Classifiers for Multilingual Pretraining Data Selection}, 
      author={Yassine Turki and Vinko Sabolčec and Bettina Messmer and Martin Jaggi},
      year={2026},
      eprint={2604.20549},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2604.20549}, 
}