olmocr/scripts/infinigram_count.py

#!/usr/bin/env python3
import argparse
import json
import random
import re
import time

import boto3
import requests
from tqdm import tqdm
from transformers import AutoTokenizer

# Allowed characters: alphanumeric, space, and basic punctuation ".,!?()"
ALLOWED_RE = re.compile(r"^[A-Za-z0-9\.,!?() ]+$")


def get_random_line_from_s3(bucket, key):
    """
    Reads an S3 object line-by-line and returns a random line using reservoir sampling.
    """
    s3 = boto3.client("s3")
    response = s3.get_object(Bucket=bucket, Key=key)
    random_line = None
    count = 0
    for line in response["Body"].iter_lines():
        if not line:
            continue
        line_str = line.decode("utf-8")
        count += 1
        if random.randint(1, count) == 1:
            random_line = line_str
    return random_line


def query_infinigram(ngram, index="v4_rpj_llama_s4", retries=3):
    """
    Sends a count query to the infini-gram API for the given n-gram.
    Retries a few times in case of network issues.
    """
    url = "https://api.infini-gram.io/"
    payload = {
        "index": index,
        "query_type": "count",
        "query": ngram,
    }
    for i in range(retries):
        try:
            response = requests.post(url, json=payload, timeout=10)
            if response.status_code == 200:
                result = response.json()
                if "count" in result:
                    return result["count"]
        except Exception as e:  # type: ignore
            time.sleep(1)
    return 0


def process_document(doc, tokenizer, ngram_size, num_samples, index="v4_rpj_llama_s4"):
    """
    Tokenizes the document using the Llama2 tokenizer and samples random n-grams.
    Each n-gram is chosen such that:
      1. It starts on a word-split boundary (using the offset mapping and a check on the preceding character).
      2. Its decoded string contains only alphanumeric characters, spaces, and the punctuation marks ".,!?()".

    Each valid n-gram is then queried using the infini-gram API.
    The function returns the document id, the number of matching n-grams (i.e. API count > 0),
    the total number of valid n-grams sampled, and a list of tuples (flag, ngram_string).
    """
    text = doc.get("text", "")
    doc_id = doc.get("id", "Unknown")
    # Get tokenized representation with offset mapping to determine word boundaries.
    tokenized = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
    token_ids = tokenized["input_ids"]
    # offsets = tokenized["offset_mapping"]

    if len(token_ids) < ngram_size:
        return doc_id, 0, 0, []

    # Determine valid starting indices based on word-split boundaries.
    valid_positions = []
    # for i in range(len(token_ids) - ngram_size + 1):
    #     start_offset = offsets[i][0]
    #     if start_offset == 0 or (start_offset > 0 and text[start_offset - 1] == " "):
    #         valid_positions.append(i)

    if not valid_positions:
        # Fallback: if no valid positions are found, use all possible positions.
        valid_positions = list(range(len(token_ids) - ngram_size + 1))

    valid_ngram_details = []
    attempts = 0
    max_attempts = num_samples * 10  # Limit to prevent infinite loops.
    while len(valid_ngram_details) < num_samples and attempts < max_attempts:
        idx = random.choice(valid_positions)
        ngram_token_ids = token_ids[idx : idx + ngram_size]
        ngram_str = tokenizer.decode(ngram_token_ids, clean_up_tokenization_spaces=True)
        # Only accept n-grams that contain only allowed characters.
        if ALLOWED_RE.fullmatch(ngram_str) and len(ngram_str.strip()) > ngram_size * 3:
            count = query_infinigram(ngram_str, index=index)
            flag = "YES" if count > 0 else "NO"
            valid_ngram_details.append((flag, ngram_str))
        attempts += 1

    match_count = sum(1 for flag, _ in valid_ngram_details if flag == "YES")
    sample_count = len(valid_ngram_details)
    return doc_id, match_count, sample_count, valid_ngram_details


def main():
    parser = argparse.ArgumentParser(description="Infini-gram n-gram matching script with Llama2 tokenization.")
    parser.add_argument("N", type=int, help="Number of random .jsonl files to process")
    parser.add_argument("s3_path", type=str, help="S3 path to a prefix containing .jsonl files (e.g., s3://my-bucket/my-prefix/)")
    parser.add_argument("--index", type=str, default="v4_dolma-v1_7_llama", help="Infini-gram index to use (default: v4_rpj_llama_s4)")
    parser.add_argument("--ngram_size", type=int, default=10, help="Size of the n-gram to sample (default: 10)")
    parser.add_argument("--num_ngrams", type=int, default=100, help="Number of random n-grams to sample from each document (default: 100)")
    args = parser.parse_args()

    if not args.s3_path.startswith("s3://"):
        print("Error: s3_path must start with 's3://'")
        return
    path_without_scheme = args.s3_path[5:]
    parts = path_without_scheme.split("/", 1)
    bucket = parts[0]
    prefix = parts[1] if len(parts) > 1 else ""

    print("Listing .jsonl files from S3...")
    s3 = boto3.client("s3")
    response = s3.list_objects_v2(Bucket=bucket, Prefix=prefix)
    files = [obj["Key"] for obj in response.get("Contents", []) if obj["Key"].endswith(".jsonl")]
    if not files:
        print("No .jsonl files found in the given prefix.")
        return

    if args.N > len(files):
        print(f"Requested {args.N} files, but only found {len(files)}. Processing all available files.")
        args.N = len(files)
    random_files = random.sample(files, args.N)

    print("Loading Llama2 tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")

    total_matches = 0
    total_ngrams_sampled = 0

    for key in tqdm(random_files, desc="Processing files"):
        line = get_random_line_from_s3(bucket, key)
        if not line:
            print(f"Skipping {key}: No valid lines found.")
            continue
        try:
            doc = json.loads(line)
        except Exception as e:
            print(f"Error parsing JSON in {key}: {e}")
            continue
        doc_id, match_count, sample_count, details = process_document(doc, tokenizer, args.ngram_size, args.num_ngrams, index=args.index)

        # Print per-document n-gram summary
        print(f"\nDocument ID: {doc_id}")
        for flag, ngram in details:
            # Print the flag in a fixed-width field (4 characters) followed by the n-gram representation.
            print(f"{flag:4} {repr(ngram)}")
        percentage = (match_count / sample_count * 100) if sample_count else 0
        print(f"Matched n-grams: {match_count}/{sample_count} ({percentage:.2f}%)")

        total_matches += match_count
        total_ngrams_sampled += sample_count

    overall_percentage = (total_matches / total_ngrams_sampled * 100) if total_ngrams_sampled else 0
    print(f"\nTotal matched n-grams: {total_matches}/{total_ngrams_sampled} ({overall_percentage:.2f}%)")


if __name__ == "__main__":
    main()