Add main and replace argparse with flags to run sharding on g3

[michelle-yooh]

[michelle-yooh]

These are the changes:

1. Adding main function as blaze requires the use of main.
2. Replacing argparse with with flags as a workaround to avoid issues in inputs
3. Re-ordering imports and adding some lints following google3 guideline

I've tested this runs on clodutop with blaze + CPU I've also run following the conventional way of using multihost_runner + cloud TPU