| """TeenageWorks batch inference CLI. Built with Qwen; see LICENSE.""" |
| import argparse |
| import json |
| from pathlib import Path |
| from toddler_jev import ToddlerJev |
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument('--input', required=True, help='JSONL: audio, question, options; optional id') |
| p.add_argument('--output', required=True) |
| p.add_argument('--batch-size', type=int, default=1) |
| p.add_argument('--model', default=str(Path(__file__).parent)) |
| p.add_argument('--revision', default='v1.0.0') |
| p.add_argument('--device', default=None) |
| args = p.parse_args() |
| if args.batch_size < 1: |
| p.error('--batch-size must be positive') |
| source, target = Path(args.input).resolve(), Path(args.output).resolve() |
| if source == target or target.exists(): |
| p.error('Output must be a new file') |
| model = ToddlerJev.from_pretrained(args.model, revision=args.revision, device=args.device) |
| def flush(batch, output): |
| for row, result in zip(batch, model.predict_batch(batch)): |
| if 'id' in row: |
| result['id'] = row['id'] |
| output.write(json.dumps(result) + '\n') |
| output.flush() |
| with source.open(encoding='utf-8') as source_file, target.open('x', encoding='utf-8') as output: |
| batch = [] |
| for line in source_file: |
| if not line.strip(): |
| continue |
| row = json.loads(line) |
| audio = Path(row['audio']) |
| row['audio'] = str(audio if audio.is_absolute() else source.parent / audio) |
| batch.append(row) |
| if len(batch) == args.batch_size: |
| flush(batch, output) |
| batch = [] |
| if batch: |
| flush(batch, output) |
|
|
| if __name__ == '__main__': |
| main() |
|
|
