Cli

Command-line helpers for normalizing values in JSON Lines files.

Functions:

Name	Description
process_json_lines_file	Normalize values in a JSON Lines file with a normalizer.
main	Run the command-line interface.

NormalizationMethod(add_arguments, create_normalizer, help, description) dataclass

Configure a normalization method for the command-line interface.

Parameters:

Name	Type	Description	Default
add_arguments	Callable[[ArgumentParser], None]	Function that adds method-specific command-line arguments.	required
create_normalizer	Callable[[Namespace], Normalizer]	Function that creates a normalizer from parsed command-line arguments.	required

process_json_lines_file(input_path, output_path, read_key, write_key, normalizer, parent_keys=None, encoding='utf-8', replace_input=False)

Normalize values in a JSON Lines file.

Parameters:

Name	Type	Description	Default
input_path	Path	Path to the input JSON Lines file.	required
output_path	Path	Path to the normalized JSON Lines file.	required
read_key	str	Key to extract the value or list of values from each JSON object.	required
write_key	str	Key to write the normalized value or list of normalized values to.	required
normalizer	Normalizer	Function used to normalize one value.	required
parent_keys	list[str] | None	Parent keys to navigate before reading a value. Each level may contain a nested dictionary or a list of nested dictionaries.	None
encoding	str	Encoding to use for reading and writing JSON Lines files.	'utf-8'
replace_input	bool	If True, the input file will be replaced with the normalized output. A backup of the original input file will be saved to output_path.	False

Source code in src/kibad_llm/normalization/cli.py

def process_json_lines_file(
    input_path: Path,
    output_path: Path,
    read_key: str,
    write_key: str,
    normalizer: Normalizer,
    parent_keys: list[str] | None = None,
    encoding: str = "utf-8",
    replace_input: bool = False,
) -> None:
    """Normalize values in a JSON Lines file.

    Args:
        input_path: Path to the input JSON Lines file.
        output_path: Path to the normalized JSON Lines file.
        read_key: Key to extract the value or list of values from each JSON object.
        write_key: Key to write the normalized value or list of normalized values to.
        normalizer: Function used to normalize one value.
        parent_keys: Parent keys to navigate before reading a value. Each level may contain a
            nested dictionary or a list of nested dictionaries.
        encoding: Encoding to use for reading and writing JSON Lines files.
        replace_input: If True, the input file will be replaced with the normalized output.
            A backup of the original input file will be saved to output_path.

    """
    logger.info(f"Normalizing values in JSON Lines file: {input_path}")
    if input_path == output_path:
        raise ValueError("Input and output paths must be different.")
    if output_path.is_file():
        logger.warning(f"Output file {output_path} already exists. It will be overwritten.")

    if replace_input:
        # in the case of when we want to replace the input file, we first save a backup of
        # the original input file to the output path and then swap the input and output paths
        # so that the normalized values are written to the original input file
        backup_path = output_path
        input_path.replace(backup_path)
        logger.info(f"Backup of the original input file was saved to: {backup_path}")
        output_path = input_path
        input_path = backup_path

    # first, collect all unique values
    with open(input_path, encoding=encoding) as input_file:
        all_values = []
        for line in input_file:
            json_obj = json.loads(line)
            values = _collect_values_json_object(json_obj, parent_keys or [], read_key)
            all_values.extend(values)

    all_unique_values = set(all_values)
    logger.info(
        f"Unique values found for key={read_key} and parent_keys={parent_keys}: {len(all_unique_values)}"
    )

    # then, normalize all unique values (show progress bar since this may take some time)
    value2normalized = {value: normalizer(value) for value in tqdm(sorted(all_unique_values))}
    logger.info(
        f"Normalized values found: {len([v for v in value2normalized.values() if v is not None])}"
    )

    # finally, process the file again and write normalized values
    num_processed, num_normalized, num_lines = 0, 0, 0
    with open(input_path, encoding=encoding) as input_file:
        with open(output_path, "w", encoding=encoding) as output_file:
            for line in input_file:
                json_obj = json.loads(line)
                new_processed, new_normalized = _process_json_object(
                    json_obj,
                    parent_keys or [],
                    read_key,
                    write_key,
                    lambda value: value2normalized.get(value),
                )
                num_processed += new_processed
                num_normalized += new_normalized
                output_file.write(json.dumps(json_obj, ensure_ascii=False) + "\n")
                num_lines += 1
    logger.info(
        f"Processed {num_lines} lines, normalized {num_normalized} values out of {num_processed} processed."
    )
    logger.info(f"Result was written to: {output_path}")

main()

Parse command-line arguments and normalize values in one or multiple JSON Lines files.

Source code in src/kibad_llm/normalization/cli.py

def main() -> None:
    """Parse command-line arguments and normalize values in one or multiple JSON Lines files."""
    parser = argparse.ArgumentParser(
        description="Normalize values in one or multiple JSON Lines files.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    method_parsers = parser.add_subparsers(
        dest="method", required=True, help="Normalization method to use."
    )
    for method_name, method in NORMALIZATION_METHODS.items():
        method_parser = method_parsers.add_parser(
            method_name,
            description=method.description,
            help=method.help,
            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        )
        _add_common_arguments(method_parser)
        method.add_arguments(method_parser)

    args = parser.parse_args()
    method = NORMALIZATION_METHODS[args.method]
    write_key = args.write_key or f"{args.method}_normalized"
    normalizer = lru_cache(maxsize=args.cache_size)(method.create_normalizer(args))

    if args.output_paths is not None:
        if len(args.input_paths) != len(args.output_paths):
            raise ValueError(
                "Number of input paths must match number of output paths if both are provided."
            )
        output_paths = args.output_paths
    else:
        if args.replace_input:
            output_paths = [
                input_path.with_name(f"{input_path.stem}_{args.read_key}_backup.jsonl")
                for input_path in args.input_paths
            ]
        else:
            output_paths = [
                input_path.with_name(f"{input_path.stem}_{args.read_key}_normalized.jsonl")
                for input_path in args.input_paths
            ]

    for input_path, output_path in zip(args.input_paths, output_paths):
        process_json_lines_file(
            input_path=input_path,
            read_key=args.read_key,
            output_path=output_path,
            write_key=write_key,
            parent_keys=args.parent_keys,
            normalizer=normalizer,
            encoding=args.encoding,
            replace_input=args.replace_input,
        )