enzyme-ec-v2 / inputs.py
Syam26's picture
Upload folder using huggingface_hub
7ca06ca verified
Raw History Blame Contribute Delete
1.03 kB
"""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}')