()
| 20 | |
| 21 | |
| 22 | def parse_args(): |
| 23 | parser = argparse.ArgumentParser() |
| 24 | parser.add_argument("--batch_size", type=int, default=8, |
| 25 | help="Number of text samples processed per GPU step.") |
| 26 | parser.add_argument("--local_rank", type=int, default=-1, |
| 27 | help="Local rank for DDP.") |
| 28 | parser.add_argument("--output_dir", type=str, default="./sonar_embeddings_ts_100_shuffle", |
| 29 | help="Where to save the output PT files.") |
| 30 | return parser.parse_args() |
| 31 | |
| 32 | |
| 33 | class TinyStoriesTextDataset(Dataset): |