almanach/benchmark-in-a-haystack
4
1import os2import re3import torch4import fasttext5from abc import abstractmethod6from rich.console import Console7from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn, TimeElapsedColumn8from tqdm import tqdm9from utils import (10 DocumentClassifier,11 score_documents,12 load_fasttext_model,13 download_fasttext_model,14 download_transformer_model15)16 17 18console = Console()19 20class DCLMClassifier(DocumentClassifier):21 """Output score between 0 and 1."""22 23 def __init__(self, classifier_config=None):24 super().__init__(classifier_config)25 console.log("[bold cyan]Initializing DCLMClassifier...[/bold cyan]")26 models_dir = classifier_config.get("models_dir", "models") if classifier_config else "models"27 self.model = self._load_model(models_dir)28 29 @staticmethod30 def download_model(models_dir="models"):31 """Download the DCLM model to the specified directory."""32 download_fasttext_model(33 hub_repo="mlfoundations/fasttext-oh-eli5",34 hub_filename="openhermes_reddit_eli5_vs_rw_v2_bigram_200k_train.bin",35 local_filename="openhermes_reddit_eli5_vs_rw_v2_bigram_200k_train.bin",36 models_dir=models_dir37 )38 39 @staticmethod40 def _load_model(models_dir="models"):41 model_path = os.path.join(models_dir, "openhermes_reddit_eli5_vs_rw_v2_bigram_200k_train.bin")42 if not os.path.exists(model_path):43 console.log(f"[yellow]Model not found at {model_path}. Downloading...[/yellow]")44 download_fasttext_model(45 hub_repo="mlfoundations/fasttext-oh-eli5",46 hub_filename="openhermes_reddit_eli5_vs_rw_v2_bigram_200k_train.bin",47 local_filename="openhermes_reddit_eli5_vs_rw_v2_bigram_200k_train.bin",48 models_dir=models_dir49 )50 return load_fasttext_model(model_path)51 52 def _score_documents(self, documents):53 console.log("[bold cyan]Scoring documents with DCLMClassifier...[/bold cyan]")54 return score_documents(documents, self.model)55 56class TextbookFastTextClassifier(DocumentClassifier):57 """Output score between 0 and 1."""58 59 def __init__(self, classifier_config=None):60 super().__init__(classifier_config)61 console.log("[bold cyan]Initializing TextbookFastTextClassifier...[/bold cyan]")62 models_dir = classifier_config.get("models_dir", "models") if classifier_config else "models"63 self.model = self._load_model(models_dir)64 65 @staticmethod66 def download_model(models_dir="models"):67 """Download the Textbook FastText model to the specified directory."""68 download_fasttext_model(69 hub_repo="kenhktsui/llm-data-textbook-quality-fasttext-classifer-v1",70 hub_filename="model.bin",71 local_filename="textbook_model.bin",72 models_dir=models_dir73 )74 75 @staticmethod76 def _load_model(models_dir="models"):77 model_path = os.path.join(models_dir, "textbook_model.bin")78 if os.path.exists(model_path):79 console.log(f"[yellow]Loading Textbook FastText model from local {model_path}...[/yellow]")80 return fasttext.load_model(model_path)81 else:82 console.log("[yellow]Model not found locally. Downloading Textbook FastText model...[/yellow]")83 download_fasttext_model(84 hub_repo="kenhktsui/llm-data-textbook-quality-fasttext-classifer-v1",85 hub_filename="model.bin",86 local_filename="textbook_model.bin",87 models_dir=models_dir88 )89 return fasttext.load_model(model_path)90 91 def _score_documents(self, documents):92 console.log("[bold cyan]Scoring documents with TextbookFastTextClassifier...[/bold cyan]")93 94 def _hq_prob(labels, scores):95 """Extract probability of the HIGH_QUALITY class."""96 for label, score in zip(labels, scores):97 if label.lstrip("__label__").upper() == "HIGH_QUALITY":98 return float(score)99 return 0.0100 101 texts = [re.sub(r"\n+", " ", doc["text"]) for doc in tqdm(documents, desc="🔄 Preprocessing text", unit="doc")]102 console.log("[yellow]Running FastText inference (C++ backend, no progress available)...[/yellow]")103 preds = self.model.predict(texts, k=2)104 results = []105 for doc, labels, scores in tqdm(zip(documents, preds[0], preds[1]), desc="📊 Formatting results", total=len(documents), unit="doc"):106 score = _hq_prob(labels, scores)107 results.append({108 "id": doc["id"],109 "source": doc["source"],110 "contains_benchmark": doc["contains_benchmark"],111 "benchmark_type": doc["benchmark_type"],112 "benchmark_index": doc.get("benchmark_index", None),113 "score": float(score)114 })115 return results116 117class TransformerClassifier(DocumentClassifier):118 119 def __init__(self, classifier_config=None):120 super().__init__(classifier_config)121 console.log(f"[bold cyan]Initializing {self.__class__.__name__}...[/bold cyan]")122 config = self.get_model_config()123 models_dir = classifier_config.get("models_dir", "models") if classifier_config else "models"124 model_dir = os.path.join(models_dir, os.path.basename(config['model_dir']))125 self.tokenizer, self.model, self.device = self._load_transformer_model(126 model_dir, 127 config['hub_name'], 128 config.get('trust_remote_code', False),129 config.get('torch_dtype')130 )131 self.batch_size = classifier_config.get('batch_size', 16) if classifier_config else 16132 133 @classmethod134 def download_model(cls, models_dir="models"):135 """Download the transformer model to the specified directory."""136 # Create a temporary instance to get config (without initializing full model)137 config = cls.__new__(cls).get_model_config()138 local_dirname = os.path.basename(config['model_dir'])139 140 download_transformer_model(141 hub_name=config['hub_name'],142 local_dirname=local_dirname,143 models_dir=models_dir,144 trust_remote_code=config.get('trust_remote_code', False),145 torch_dtype=config.get('torch_dtype')146 )147 148 @abstractmethod149 def get_model_config(self):150 pass151 152 @abstractmethod153 def process_outputs(self, outputs, doc_batch):154 pass155 156 def _score_documents(self, documents):157 console.log(f"[bold cyan]Scoring documents with {self.__class__.__name__}...[/bold cyan]")158 results = []159 num_batches = (len(documents) + self.batch_size - 1) // self.batch_size160 for idx_batch in tqdm(range(0, len(documents), self.batch_size), desc=f"⚡ {self.__class__.__name__}: Inference", total=num_batches, unit="batch"):161 doc_batch = documents[idx_batch:idx_batch + self.batch_size]162 text_batch = [doc["text"] for doc in doc_batch]163 164 config = self.get_model_config()165 tokenizer_kwargs = {"return_tensors": "pt", "padding": "longest", "truncation": True}166 if config.get('max_length'):167 tokenizer_kwargs["max_length"] = config['max_length']168 169 inputs = self.tokenizer(text_batch, **tokenizer_kwargs).to(self.device)170 inputs = self._process_inputs(inputs)171 172 with torch.no_grad():173 outputs = self.model(**inputs)174 175 results.extend(self.process_outputs(outputs, doc_batch))176 177 return results178 179 def _process_inputs(self, inputs):180 return inputs181 182 183class FinewebEduClassifier(TransformerClassifier):184 """Output score between 0 and 5."""185 186 def get_model_config(self):187 return {188 'model_dir': "models/fineweb-edu-classifier",189 'hub_name': "HuggingFaceTB/fineweb-edu-classifier",190 'trust_remote_code': False191 }192 193 def process_outputs(self, outputs, doc_batch):194 results = []195 for i_doc, doc in enumerate(doc_batch):196 logits = outputs.logits[i_doc].float().detach().cpu().numpy()197 score = logits.item()198 score = max(0, min(score, 5))199 int_score = int(round(score))200 results.append({201 "id": doc["id"],202 "source": doc["source"],203 "contains_benchmark": doc["contains_benchmark"],204 "benchmark_type": doc["benchmark_type"],205 "benchmark_index": doc.get("benchmark_index", None),206 "score": float(score),207 "int_score": int_score208 })209 return results210 211 212class GaperonClassifier(TransformerClassifier):213 """Output score between 0 and 1.5."""214 215 def get_model_config(self):216 return {217 'model_dir': "models/gaperon-quality-cls",218 'hub_name': "almanach/gaperon-quality-cls",219 'trust_remote_code': True,220 'max_length': 512221 }222 223 def _process_inputs(self, inputs):224 return {k: v[:, :512] for k, v in inputs.items()}225 226 def process_outputs(self, outputs, doc_batch):227 results = []228 for i_doc, doc in enumerate(doc_batch):229 logits = outputs.logits_list[0][i_doc].squeeze(0).float().softmax(-1).detach().cpu().numpy()230 score = (logits[0] + 0.5 * logits[2]).item()231 int_score = int(round(max(0, min(1+2*score, 3))))232 results.append({233 "id": doc["id"],234 "source": doc["source"],235 "contains_benchmark": doc["contains_benchmark"],236 "benchmark_type": doc["benchmark_type"],237 "benchmark_index": doc.get("benchmark_index", None),238 "score": float(score),239 "int_score": int_score240 })241 return results242 243 244class NemoCuratorEduClassifier(TransformerClassifier):245 """Output score between 0 and 5."""246 247 def get_model_config(self):248 return {249 'model_dir': "models/nemocurator-fineweb-mixtral-edu-classifier",250 'hub_name': "nvidia/nemocurator-fineweb-mixtral-edu-classifier",251 'trust_remote_code': False,252 'max_length': 512,253 'torch_dtype': torch.bfloat16254 }255 256 def process_outputs(self, outputs, doc_batch):257 results = []258 for i_doc, doc in enumerate(doc_batch):259 logit = outputs.logits[i_doc].squeeze(-1).float().cpu().numpy()260 score = float(logit)261 score = max(0, min(score, 5))262 int_score = int(round(score))263 pred_label = "high_quality" if score >= 2.5 else "low_quality"264 results.append({265 "id": doc["id"],266 "source": doc["source"],267 "contains_benchmark": doc["contains_benchmark"],268 "benchmark_type": doc["benchmark_type"],269 "benchmark_index": doc.get("benchmark_index", None),270 "score": score,271 "int_score": int_score,272 "label": pred_label273 })274 return results275 276 277class FinePDFsClassifierBase(DocumentClassifier):278 """Output score: unbounded."""279 280 def __init__(self, classifier_config=None):281 super().__init__(classifier_config)282 console.log(f"[bold cyan]Initializing {self.__class__.__name__}...[/bold cyan]")283 config = self.get_model_config()284 models_dir = classifier_config.get("models_dir", "models") if classifier_config else "models"285 model_dir = os.path.join(models_dir, os.path.basename(config['model_dir']))286 self.tokenizer, self.model, self.device = self._load_transformer_model(287 model_dir, config['hub_name']288 )289 self.CHUNK_SIZE = 2046290 self.MAX_CHARS = 10_000291 self.batch_size = classifier_config.get('batch_size', 1) if classifier_config else 1292 293 @classmethod294 def download_model(cls, models_dir="models"):295 """Download the FinePDFs model to the specified directory."""296 # Create a temporary instance to get config (without initializing full model)297 config = cls.__new__(cls).get_model_config()298 local_dirname = os.path.basename(config['model_dir'])299 300 download_transformer_model(301 hub_name=config['hub_name'],302 local_dirname=local_dirname,303 models_dir=models_dir304 )305 306 @abstractmethod307 def get_model_config(self):308 pass309 310 def _trim_to_whitespace(self, text, trim_start, trim_end):311 if trim_start:312 match = re.search(r'\s', text)313 text = text[match.start()+1:] if match else text[10:]314 if trim_end:315 match = re.search(r'\s', text[::-1])316 text = text[:len(text) - match.start() - 1] if match else text[:-10]317 return text318 319 def _create_text_chunks(self, text):320 if len(text) <= 2 * self.MAX_CHARS:321 tokens = self.tokenizer.encode(text[:self.MAX_CHARS], return_tensors="np", add_special_tokens=False)[0]322 chunk_text = self.tokenizer.decode(tokens[:self.CHUNK_SIZE], skip_special_tokens=True)323 return [self._trim_to_whitespace(chunk_text, False, True)]324 325 text_top, text_bottom = text[:self.MAX_CHARS], text[-self.MAX_CHARS:]326 tokens = self.tokenizer.batch_encode_plus([text_top, text_bottom], return_tensors="np", add_special_tokens=False)["input_ids"]327 chunks = [tokens[0][:self.CHUNK_SIZE], tokens[1][-self.CHUNK_SIZE:]]328 chunks_text = self.tokenizer.batch_decode(chunks, skip_special_tokens=True)329 return [330 self._trim_to_whitespace(chunks_text[0], False, True),331 self._trim_to_whitespace(chunks_text[1], True, False)332 ]333 334 def _score_documents(self, documents):335 console.log(f"[bold cyan]Scoring documents with {self.__class__.__name__}...[/bold cyan]")336 results = []337 num_batches = (len(documents) + self.batch_size - 1) // self.batch_size338 339 for idx_batch in tqdm(range(0, len(documents), self.batch_size), desc=f"⚡ {self.__class__.__name__}: Inference", total=num_batches, unit="batch"):340 doc_batch = documents[idx_batch:idx_batch + self.batch_size]341 342 all_chunks = []343 doc_chunk_mapping = []344 345 for doc_idx, doc in enumerate(doc_batch):346 chunks = self._create_text_chunks(doc["text"])347 chunk_start_idx = len(all_chunks)348 all_chunks.extend(chunks)349 doc_chunk_mapping.append((doc_idx, chunk_start_idx, len(all_chunks)))350 351 if all_chunks:352 inputs = self.tokenizer(all_chunks, return_tensors="pt", padding="longest", truncation=True).to(self.device)353 with torch.no_grad():354 outputs = self.model(**inputs)355 all_scores = outputs.logits.squeeze(-1).float().detach().cpu().numpy()356 357 if len(all_chunks) == 1:358 all_scores = [all_scores.item()]359 else:360 all_scores = all_scores.tolist()361 362 for doc_idx, chunk_start, chunk_end in doc_chunk_mapping:363 doc = doc_batch[doc_idx]364 doc_scores = all_scores[chunk_start:chunk_end]365 final_score = max(doc_scores)366 367 results.append({368 "id": doc["id"],369 "source": doc["source"],370 "contains_benchmark": doc["contains_benchmark"],371 "benchmark_type": doc["benchmark_type"],372 "benchmark_index": doc.get("benchmark_index", None),373 "score": float(final_score),374 "int_score": int(round(max(0, min(final_score, 5))))375 })376 377 return results378 379 380class FinePDFsEduClassifier(FinePDFsClassifierBase):381 """Output score: unbounded."""382 383 def get_model_config(self):384 return {385 'model_dir': "models/finepdfs-edu-classifier-eng-Latn",386 'hub_name': "HuggingFaceFW/finepdfs_edu_classifier_eng_Latn"387 }388 389 390class FinePDFsEduClassifierV2(FinePDFsClassifierBase):391 """Output score: unbounded."""392 393 def get_model_config(self):394 return {395 'model_dir': "models/finepdfs-edu-classifier-v2-eng-Latn",396 'hub_name': "HuggingFaceFW/finepdfs_edu_classifier_v2_eng_Latn"397 }398 399 400class FinePDFsDCLMClassifier(FinePDFsClassifierBase):401 """Output score: unbounded."""402 403 def get_model_config(self):404 return {405 'model_dir': "models/finepdfs-dclm-classifier-eng-Latn",406 'hub_name': "HuggingFaceFW/finepdfs_dclm_classifier_eng_Latn"407 }408 409 410class EuroFilterClassifier(TransformerClassifier):411 """Output score between 0 and 5."""412 413 def get_model_config(self):414 return {415 'model_dir': "models/eurofilter-v1",416 'hub_name': "utter-project/EuroFilter-v1",417 'trust_remote_code': True,418 'max_length': 512,419 'torch_dtype': torch.bfloat16420 }421 422 def process_outputs(self, outputs, doc_batch):423 results = []424 for i_doc, doc in enumerate(doc_batch):425 score = outputs.logits[i_doc].squeeze().float().cpu().numpy().item()426 score = max(0, min(score, 5))427 int_score = int(round(score))428 429 prob = torch.nn.functional.sigmoid(outputs.binary_logits[i_doc]).float().cpu().numpy().item()430 binary_pred = 1 if prob >= 0.5 else 0431 432 results.append({433 "id": doc["id"],434 "source": doc["source"],435 "contains_benchmark": doc["contains_benchmark"],436 "benchmark_type": doc["benchmark_type"],437 "benchmark_index": doc.get("benchmark_index", None),438 "score": float(score),439 "int_score": int_score,440 "binary_pred": binary_pred,441 "prob": float(prob)442 })443 return results444