File size: 1,694 Bytes
c83c617 | 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 | """Evaluate a C1 checkpoint only on the frozen nonblind snapshots."""
import argparse
from pathlib import Path
import sys
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
def main():
from vimeml.posttraining.infer import (TRACKS, read_snapshot, load_conditional_checkpoint,
fp32_evaluation, evaluate_snapshot)
from vimeml.training.data import write_json
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
parser.add_argument("--tracks", choices=tuple(TRACKS), nargs="+", default=["standard-development"])
parser.add_argument("--condition", choices=("context_reading", "empty_reading", "empty_context"), default="context_reading")
args = parser.parse_args()
if args.output.exists() and any(args.output.iterdir()):
parser.error("Use a new evaluation output directory")
rows = {track: read_snapshot(track) for track in args.tracks}
model, codec, metadata = load_conditional_checkpoint(args.checkpoint,
ROOT / "artifacts/tokenizers/v4-a/ja-unigram-8192/tokenizer.model", args.device)
report = {"checkpoint": metadata, "condition": args.condition, "tracks": {}}
with fp32_evaluation(model):
for track, data in rows.items():
report["tracks"][track] = evaluate_snapshot(model, codec, data, args.output / track, args.condition)
write_json(args.output / "summary.json", report)
print(report)
if __name__ == "__main__":
main()
|