mirror of https://github.com/razor-ai/soup.git
203 lines
6.5 KiB
Python
203 lines
6.5 KiB
Python
"""Data quality scoring — perplexity and coherence filters."""
|
||
|
||
import logging
|
||
import math
|
||
from typing import Optional
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def compute_perplexity_scores(
|
||
texts: list[str],
|
||
model_name: str = "gpt2",
|
||
batch_size: int = 8,
|
||
max_length: int = 512,
|
||
) -> list[float]:
|
||
"""Compute perplexity scores for a list of texts using a language model.
|
||
|
||
Lower perplexity = more "normal" text (the model finds it predictable).
|
||
Very high perplexity may indicate garbage, corrupted, or very unusual text.
|
||
|
||
Args:
|
||
texts: List of text strings to score.
|
||
model_name: HuggingFace model to use for perplexity (default: gpt2).
|
||
batch_size: Batch size for inference.
|
||
max_length: Max tokens per text.
|
||
|
||
Returns:
|
||
List of perplexity scores (one per text).
|
||
"""
|
||
import torch
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||
if tokenizer.pad_token is None:
|
||
tokenizer.pad_token = tokenizer.eos_token
|
||
|
||
model = AutoModelForCausalLM.from_pretrained(model_name)
|
||
model.eval()
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
model = model.to(device)
|
||
|
||
scores = []
|
||
|
||
for start_idx in range(0, len(texts), batch_size):
|
||
batch_texts = texts[start_idx:start_idx + batch_size]
|
||
|
||
encodings = tokenizer(
|
||
batch_texts,
|
||
return_tensors="pt",
|
||
padding=True,
|
||
truncation=True,
|
||
max_length=max_length,
|
||
)
|
||
input_ids = encodings["input_ids"].to(device)
|
||
attention_mask = encodings["attention_mask"].to(device)
|
||
|
||
with torch.no_grad():
|
||
outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
|
||
# Per-token loss: outputs.loss is mean over all tokens
|
||
# We need per-sample loss for individual perplexity scores
|
||
shift_logits = outputs.logits[:, :-1, :].contiguous()
|
||
shift_labels = input_ids[:, 1:].contiguous()
|
||
shift_mask = attention_mask[:, 1:].contiguous()
|
||
|
||
# Set pad positions to -100 so CrossEntropyLoss skips them
|
||
shift_labels = shift_labels.masked_fill(shift_mask == 0, -100)
|
||
|
||
loss_fct = torch.nn.CrossEntropyLoss(reduction="none", ignore_index=-100)
|
||
per_token_loss = loss_fct(
|
||
shift_logits.view(-1, shift_logits.size(-1)),
|
||
shift_labels.view(-1),
|
||
).view(shift_labels.size())
|
||
|
||
# Mask padding tokens and compute per-sample mean loss
|
||
masked_loss = per_token_loss * shift_mask.float()
|
||
sample_lengths = shift_mask.sum(dim=1).float().clamp(min=1)
|
||
sample_loss = masked_loss.sum(dim=1) / sample_lengths
|
||
|
||
for loss_val in sample_loss:
|
||
ppl = math.exp(min(loss_val.item(), 100)) # cap to avoid overflow
|
||
scores.append(ppl)
|
||
|
||
# Cleanup GPU memory
|
||
del model
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
return scores
|
||
|
||
|
||
def compute_coherence_score(texts: list[str]) -> list[float]:
|
||
"""Compute coherence scores for a list of texts.
|
||
|
||
Coherence is measured by:
|
||
- Sentence-level structure (has proper sentences)
|
||
- Word repetition ratio (low repetition = more coherent)
|
||
- Character diversity (not just repeated chars)
|
||
|
||
Returns scores in [0, 1] where 1 = highly coherent.
|
||
"""
|
||
scores = []
|
||
for text in texts:
|
||
score = _score_coherence(text)
|
||
scores.append(score)
|
||
return scores
|
||
|
||
|
||
def _score_coherence(text: str) -> float:
|
||
"""Score a single text for coherence. Returns 0.0–1.0."""
|
||
if not text or not text.strip():
|
||
return 0.0
|
||
|
||
text = text.strip()
|
||
words = text.split()
|
||
|
||
if len(words) < 3:
|
||
return 0.2
|
||
|
||
# 1. Word repetition ratio (0-1, lower repetition = higher score)
|
||
unique_words = set(w.lower() for w in words)
|
||
repetition_score = len(unique_words) / len(words) # 1.0 = all unique
|
||
|
||
# 2. Sentence structure: check for sentence-ending punctuation
|
||
sentence_endings = sum(1 for ch in text if ch in ".!?")
|
||
sentence_score = min(1.0, sentence_endings / max(1, len(words) / 15))
|
||
|
||
# 3. Character diversity: ratio of unique chars to total
|
||
unique_chars = set(text.lower())
|
||
char_diversity = min(1.0, len(unique_chars) / 30) # 30+ unique chars = 1.0
|
||
|
||
# 4. Average word length (very short or very long = less coherent)
|
||
avg_word_len = sum(len(w) for w in words) / len(words)
|
||
word_len_score = 1.0 if 3 <= avg_word_len <= 12 else 0.5
|
||
|
||
# Weighted combination
|
||
coherence = (
|
||
0.35 * repetition_score
|
||
+ 0.25 * sentence_score
|
||
+ 0.20 * char_diversity
|
||
+ 0.20 * word_len_score
|
||
)
|
||
return round(min(1.0, max(0.0, coherence)), 4)
|
||
|
||
|
||
def filter_by_quality(
|
||
data: list[dict],
|
||
perplexity_threshold: Optional[float] = None,
|
||
coherence_threshold: Optional[float] = None,
|
||
perplexity_model: str = "gpt2",
|
||
text_field: Optional[str] = None,
|
||
) -> tuple[list[dict], list[dict]]:
|
||
"""Filter dataset rows by quality scores.
|
||
|
||
Args:
|
||
data: List of data rows (dicts).
|
||
perplexity_threshold: Max perplexity allowed (rows above this are removed).
|
||
coherence_threshold: Min coherence required (rows below this are removed).
|
||
perplexity_model: Model for perplexity scoring.
|
||
text_field: Which field to score. None = concatenate all string values.
|
||
|
||
Returns:
|
||
(kept, removed) tuple of data lists.
|
||
"""
|
||
if not data:
|
||
return [], []
|
||
|
||
# Extract texts
|
||
texts = []
|
||
for row in data:
|
||
if text_field and text_field in row:
|
||
texts.append(str(row[text_field]))
|
||
else:
|
||
texts.append(" ".join(str(v) for v in row.values() if v))
|
||
|
||
# Compute scores
|
||
perplexity_scores = None
|
||
if perplexity_threshold is not None:
|
||
perplexity_scores = compute_perplexity_scores(
|
||
texts, model_name=perplexity_model,
|
||
)
|
||
|
||
coherence_scores = None
|
||
if coherence_threshold is not None:
|
||
coherence_scores = compute_coherence_score(texts)
|
||
|
||
# Filter
|
||
kept = []
|
||
removed = []
|
||
for idx, row in enumerate(data):
|
||
remove = False
|
||
if perplexity_scores is not None and perplexity_scores[idx] > perplexity_threshold:
|
||
remove = True
|
||
if coherence_scores is not None and coherence_scores[idx] < coherence_threshold:
|
||
remove = True
|
||
|
||
if remove:
|
||
removed.append(row)
|
||
else:
|
||
kept.append(row)
|
||
|
||
return kept, removed
|