ai: start creating initial scaffolding

This commit is contained in:
Starbeamrainbowlabs 2022-07-25 19:00:53 +01:00
parent 1ec502daea
commit dac6919fcd
Signed by: sbrl
GPG key ID: 1BE5172E637709C2
2 changed files with 74 additions and 0 deletions

View file

@ -0,0 +1,33 @@
#!/usr/bin/env python3
import importlib
import sys
from loguru import logger
def init_logging():
pass
from parse_args import parse_args
def main():
subcommand, args = parse_args()
if args == None:
return
imported_module = importlib.import_module(f"subcommands.{subcommand}")
# TODO: Support multiple subcommands here
match subcommand:
case "pretrain":
imported_module.pretrain(args)
case _:
sys.stderr.write(f"Error: The subcommand '{subcommand}' hasn't been registered in index.py yet.\nThis is a bug.\n")
exit(1)
if __name__ == "__main__":
main()
else:
print("This script must be run directly. It cannot be imported.")
exit(1)

41
aimodel/src/parse_args.py Normal file
View file

@ -0,0 +1,41 @@
import argparse
import sys
import re
import importlib
# import pysnooper
# @pysnooper.snoop()
def parse_args():
"""Defines and parses the CLI arguments."""
if len(sys.argv) < 2:
sys.stderr.write("""
This program trains, manipulates, visualises, and manages a contrastive learning based rainfall radar water depth prediction model.
It functions by first finding relationships between the rainfall radar data and the water depth + heightmap data (the 'pretrain' subcommand). After this, a decoder model to predict water depth (modelled as an image segmentation task), can then be trained.
Available subcommands:
pretrain Pretrain a contrastive learning model as an encoder.
For more information, do src/index.py <subcommand> --help.
""")
exit(0)
subcommand = re.sub(r'[^a-z0-9-]', '', sys.argv[1])
subcommand_argparser = importlib.import_module(f"subcommands.{subcommand}").parse_args
parser = subcommand_argparser()
# sys.stderr.write(f"Error: Unknown subcommand '{subcommand} (try --help).\n")
# exit(1)
if parser == None:
sys.stderr.write(f"Error: The subcommand '{subcommand}' did not return an argument parser. This is a bug.\n")
exit(1)
parser.add_argument("--only-gpu",
help="If the GPU is not available, exit with an error (useful on shared HPC systems to avoid running out of memory & affecting other users)", action="store_true")
return subcommand, parser.parse_args(args=sys.argv[2:])