scripts/local_scripts/infer_folder.py

import argparse
import json
import re
import os
import unicodedata
from typing import Tuple, List
from multiprocessing import Pool

import fasttext
import pandas as pd
from tqdm import tqdm
from transformers import LlamaTokenizerFast


language_model_map = {
    "en": "classifiers/ultra_fineweb_en.bin",
    "zh": "classifiers/ultra_fineweb_zh.bin"
}

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--language", type=str, required=True, help="Inference language, support: en, zh.")
    parser.add_argument("--data-path", type=str, required=True, help="Data path.")
    parser.add_argument("--save-path", type=str, required=True, help="Save path root.")
    parser.add_argument("--content-key", type=str, required=True, help="Content key for inference.")
    parser.add_argument("--tokenizer-path", type=str, default="local_tokenizer", help="Tokenizer path.")
    parser.add_argument("--processes-num", type=int, default=64, help="Number of processes.")
    parser.add_argument("--write-batch-size", type=int, default=100, help="Write batch size.")
    parser.add_argument("--inplace", action="store_true", help="Inplace already processed data.")
    return parser.parse_args()


def fasttext_preprocess_func(content: str, tokenizer: LlamaTokenizerFast) -> str:
    """Fasttext preprocess function.

    Args:
        content (str): Content to process.

    Returns:
        str: Processed normalized content.
    """

    # 1. remove multiple newlines
    content = re.sub(r'\n{3,}', '\n\n', content)

    # 2. lower the content
    content = content.lower()

    # 3. remove diacritics
    content = ''.join(
        c for c in unicodedata.normalize('NFKD', content)
        if unicodedata.category(c) != 'Mn')

    # 4. word segmentation
    token_ids = tokenizer.encode(content, add_special_tokens=False)
    single_text_list = []
    for token_id in token_ids:
        curr_text = tokenizer.decode([token_id])
        single_text_list.append(curr_text)

    content = ' '.join(single_text_list)

    # 5. keep escape chars, \n, \t, \r -> \\n, \\t, \\r,
    # which will saved as \n, \t, \r in txt file.
    content = re.sub(r'\n', '\\\\n', content)
    content = re.sub(r'\r', '\\\\r', content)
    content = re.sub(r'\t', '\\\\t', content)
    content = re.sub(r' +', ' ', content)
    content = content.strip()

    return content


def fasttext_infer(norm_content: str, fasttext_model: fasttext.FastText) -> Tuple[str, float]:
    """Fasttext inference function

    Args:
        content (str): input text
    
    Returns:
        str: json string with pred_label and pred_score
    """

    pred_label, pred_prob = fasttext_model.predict(norm_content)
    pred_label = pred_label[0]
    _score = min(pred_prob.tolist()[0], 1)
    if pred_label == "__label__neg":
        _score = 1 - _score

    return pred_label, _score


def load_data(file_path: str, content_key: str) -> List[str]:
    """Load data from file path.
    
    Args:
        file_path (str): File path.
        content_key (str): Content key.
    
    Returns:
        List[str]: List of content.
    """
    samples = []
    if file_path.endswith(".jsonl") or file_path.endswith(".json"):
        with open(file_path, "r", encoding="utf-8") as f:
            for line in f:
                data = json.loads(line.strip())
                if content_key in data:
                    if data[content_key] == "":
                        print("Empty text, continue")
                        continue
                    if data[content_key] is None:
                        print("None text, continue")
                        continue
                    samples.append(data[content_key])
    elif file_path.endswith(".parquet"):
        df = pd.read_parquet(file_path)
        for _, row in df.iterrows():
            if content_key in row:
                if row[content_key] == "":
                    print("Empty text, continue")
                    continue
                if row[content_key] is None:
                    print("None text, continue")
                    continue
                samples.append(row[content_key])
    else:
        raise ValueError(f"Unsupported file type: {file_path}")
    return samples


def process_file(
        file_path: str,
        tokenizer_path: str,
        fasttext_model_path: str,
        save_path: str,
        item: int,
        content_key: str,
        inplace: bool,
        write_batch_size: int) -> None:
    """Process a single file.
    
    Args:
        file_path (str): File path to process.
        tokenizer_path (str): Tokenizer path.
        fasttext_model_path (str): Fasttext model path.
        save_path (str): Save path.
        item (int): Current process item index.
        content_key (str): Content key.
        write_batch_size (int): Write batch size.
    """

    # load tokenizer and fasttext model
    tokenizer = LlamaTokenizerFast.from_pretrained(tokenizer_path)
    fasttext_model = fasttext.load_model(fasttext_model_path)

    # load data
    all_texts = load_data(file_path, content_key)

    # get file name
    file_name = os.path.basename(file_path)
    curr_file_name = ".".join(file_name.split(".")[:-1])

    output_file = f"{curr_file_name}_fasttext_pos.jsonl"
    output_file = os.path.join(save_path, output_file)

    if inplace and os.path.exists(output_file):
        print(f"File {output_file} already exists, skip")
        return
    
    if os.path.exists(output_file):
        # remove the file
        print(f"File {output_file} already exists, remove it")
        os.remove(output_file)

    results = []
    print(f"ID: {item}, Begin to process {file_path}, total {len(all_texts)} samples, results will be saved in {output_file}")
    for text in tqdm(all_texts):
        norm_content = fasttext_preprocess_func(text, tokenizer)
        label, score = fasttext_infer(norm_content, fasttext_model)

        # label is __label__pos or __label__neg
        if label == "__label__pos":
            curr_result = {"content": text, "pred_label": label, "pred_score": score}
            results.append(curr_result)

        if len(results) >= write_batch_size:
            with open(output_file, "a", encoding="utf-8") as f:
                f.write("\n".join(json.dumps(r, ensure_ascii=False) for r in results) + "\n")
            results.clear()

    # process remaining results
    if results:
        with open(output_file, "a", encoding="utf-8") as f:
            f.write("\n".join(json.dumps(r, ensure_ascii=False) for r in results) + "\n")


def main():
    args = parse_args()
    language = args.language
    data_path = args.data_path
    save_path = args.save_path
    content_key = args.content_key
    tokenizer_path = args.tokenizer_path
    processes_num = args.processes_num
    write_batch_size = args.write_batch_size
    inplace = args.inplace

    assert os.path.exists(data_path), f"Data path {data_path} not exists"
    assert os.path.exists(tokenizer_path), f"Tokenizer path {tokenizer_path} not exists"

    assert language in language_model_map, f"Language {language} not supported"
    fasttext_model_path = language_model_map[language]

    if not os.path.exists(save_path):
        os.makedirs(save_path, exist_ok=True)
    
    data_path_list = os.listdir(data_path)
    data_path_list = [os.path.join(data_path, file_name) for file_name in data_path_list]

    print("=" * 100)
    print(f"Begin processing\n"
          f"- data path: {data_path}\n"
          f"- save path: {save_path}\n"
          f"- content key: {content_key}\n"
          f"- tokenizer path: {tokenizer_path}\n"
          f"- processes num: {processes_num}\n"
          f"- write batch size: {write_batch_size}\n"
          f"- inplace: {inplace}")
    print("=" * 100)

    print(f"Total {len(data_path_list)} files to process")

    # process data
    with Pool(processes=processes_num) as pool:
        pool.starmap(process_file, [(
            file_path, tokenizer_path, fasttext_model_path, save_path, item, content_key, inplace, write_batch_size)
            for item, file_path in enumerate(data_path_list)])

    print("Finished processing all files")

if __name__ == "__main__":
    main()