Mentors4EDU's picture
Publish training pipeline under training/; update repository notes
db3982a verified
Raw History Blame Contribute Delete
1.14 kB
from __future__ import annotations
import argparse
import json
import logging
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from riemann.data import prepare_all, DATA_DIR
def main() -> None:
logging.basicConfig(
level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
)
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--only",
nargs="*",
default=None,
help="only run these loader names, e.g. --only fable",
)
parser.add_argument(
"--append",
action="store_true",
help="append to existing corpus files instead of overwriting them",
)
parser.add_argument(
"--data-dir",
default=str(DATA_DIR),
help="where to write pretrain_corpus.txt and ntk_finetune.jsonl",
)
args = parser.parse_args()
result = prepare_all(
data_dir=Path(args.data_dir), only=args.only, append=args.append
)
print(json.dumps(result, indent=2))
if __name__ == "__main__":
main()