pararater

ParaRater: Enhancing Cross-Lingual Transfer in LLMs with Meta-Learning

ParaRater is a data selection method that enhances cross-lingual transfer by selecting the most valuable parallel pairs, forming high-impact parallel corpora with two meta-learned raters.

Raters

This is the repository of trained rater models of ParaRater. Raters are trained based on Qwen3-Embedding-0.6B.

Usage

Each pair of Rater1 and Rater2 trained for a specific target language can be jointly used to filter English corpora.

import argparse
import torch
import pandas as pd
import pyarrow.parquet as pq
from transformers import AutoTokenizer, AutoModelForSequenceClassification

def percentile_ranks(scores):
    # higher=better -> percentile in [0,1], 1.0 is best
    order = torch.argsort(scores, descending=True)
    ranks = torch.empty_like(order, dtype=torch.float)
    ranks[order] = torch.arange(len(scores), dtype=torch.float)
    denom = max(1, len(scores) - 1)
    return 1.0 - ranks / denom

@torch.no_grad()
def batched_logits(texts, tokenizer, model, batch_size=64, max_length=512, device="cuda" if torch.cuda.is_available() else "cpu"):
    model.to(device).eval()
    out_scores = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        enc = tokenizer(batch, padding=True, truncation=True,
                        max_length=max_length, return_tensors="pt").to(device)
        logits = model(**enc).logits.squeeze(-1)  # (B,) for class_num=1
        out_scores.append(logits.cpu())
    return torch.cat(out_scores, dim=0)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--parquet", required=True, help="Input parquet path with column 'text'.")
    ap.add_argument("--rater1", default="pararater_rater1_en-ar", help="Rater1.")
    ap.add_argument("--rater2", default="pararater_rater2_en-ar", help="Rater2.")
    ap.add_argument("--save_parquet", default=None, help="Optional output parquet for kept samples.")
    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--max_length", type=int, default=512)
    args = ap.parse_args()

    # 1) Load data
    df = pq.read_table(args.parquet).to_pandas()
    assert "text" in df.columns, "Parquet must have a 'text' column."
    texts = df["text"].astype(str).tolist()

    # 2) Load raters
    tok = AutoTokenizer.from_pretrained(args.rater1, trust_remote_code=True)
    r1 = AutoModelForSequenceClassification.from_pretrained(args.rater1, trust_remote_code=True)
    r2 = AutoModelForSequenceClassification.from_pretrained(args.rater2, trust_remote_code=True)

    # 3) Score -> percentile ranks
    s1 = batched_logits(texts, tok, r1, batch_size=args.batch_size, max_length=args.max_length)
    s2 = batched_logits(texts, tok, r2, batch_size=args.batch_size, max_length=args.max_length)
    p1 = percentile_ranks(s1)  # 1.0 best
    p2 = percentile_ranks(s2)

    # 4) Rule: keep if (p1 >= 0.6) and (p2 <= p1 - 0.2)
    top = p1 >= 0.6
    drop = p2 <= (p1 - 0.2)
    keep_mask = (top & drop).numpy()

    kept = df.loc[keep_mask]
    print(f"Total: {len(df)} | Rater1 top-0.6: {int(top.sum().item())} | Kept(final): {keep_mask.sum()}")

    if args.save_parquet:
        kept.to_parquet(args.save_parquet, index=False)
        print(f"Saved: {args.save_parquet}")

if __name__ == "__main__":
    main()