diff --git a/locimend/main.py b/locimend/main.py index 387cbb5..ad9e787 100644 --- a/locimend/main.py +++ b/locimend/main.py @@ -1,7 +1,8 @@ from asyncio import run from argparse import ArgumentParser, Namespace +from time import time -from model import infer_sequence, train_model +from locimend.model import infer_sequence, train_model def parse_arguments() -> Namespace: @@ -21,7 +22,10 @@ def parse_arguments() -> Namespace: async def execute_task(args): if args.task == "train": + start_time = time() train_model(data_file=args.data_file, label_file=args.label_file) + end_time = time() + print(f"Training time: {end_time - start_time}") else: prediction = await infer_sequence(sequence=args.sequence) print(f"Error-corrected sequence: {prediction}")