File size: 1,454 Bytes
e4cdf3c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """Train a small EC-frequency baseline using only the training split."""
import argparse
from collections import Counter
import json
from pathlib import Path
import pandas as pd
def train(path, output):
root = Path(__file__).parent
task = json.loads((root / 'task.json').read_text())
valid = set(json.loads((root / 'ec_vocabulary.json').read_text()))
frame = pd.read_parquet(path)
unique = frame.drop_duplicates(task['input_columns'] + ['ec_number'])
if not set(unique.ec_number) <= valid:
raise ValueError('Training data contains ECs outside the pinned vocabulary.')
counts = Counter(unique.ec_number)
ranked = sorted(counts, key=lambda ec: (-counts[ec], ec))
if not ranked:
raise ValueError('Training data is empty.')
limit = 1 # Both tasks are scored on the first EC only.
result = {'prediction': ranked[:limit], 'training_pairs': len(unique),
'baseline': 'EC frequency; ignores the input'}
target = Path(output)
target.parent.mkdir(parents=True, exist_ok=True)
target.write_text(json.dumps(result, indent=2) + '\n')
print(f'Trained on {len(unique)} input–EC pairs; saved {target}')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--train', required=True)
parser.add_argument('--output', default='weights/baseline.json')
args = parser.parse_args()
train(args.train, args.output)
|