Blame for inputs.py

"""Prepare exactly the unique input rows used by the evaluator."""
import argparse
import json
from pathlib import Path

import pandas as pd


def prepare_inputs(frame):
    config = json.loads((Path(__file__).parent / 'task.json').read_text())
    columns = ['id'] + config['input_columns']
    if not set(columns) <= set(frame.columns):
        raise ValueError('Missing input columns: ' + ', '.join(columns))
    if frame.id.isna().any() or frame.id.duplicated().any():
        raise ValueError('Input IDs must be present and unique.')
    return frame[columns].drop_duplicates(config['input_columns'], keep='first').reset_index(drop=True)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--input', required=True)
    parser.add_argument('--output', required=True)
    args = parser.parse_args()
    frame = prepare_inputs(pd.read_parquet(args.input))
    frame.to_parquet(args.output, index=False)
    print(f'Wrote {len(frame)} unique inputs to {args.output}')