| from transformers import BartTokenizer | |
| from idiomify.data import IdiomifyDataModule | |
| CONFIG = { | |
| "literal2idiomatic_ver": "pie_v0", | |
| "batch_size": 20, | |
| "num_workers": 4, | |
| "shuffle": True | |
| } | |
| def main(): | |
| tokenizer = BartTokenizer.from_pretrained("facebook/bart-large") | |
| datamodule = IdiomifyDataModule(CONFIG, tokenizer) | |
| datamodule.prepare_data() | |
| datamodule.setup() | |
| for batch in datamodule.train_dataloader(): | |
| srcs, tgts_r, tgts = batch | |
| print(srcs.shape) | |
| print(tgts_r.shape) | |
| print(tgts.shape) | |
| if __name__ == '__main__': | |
| main() | |