mahalrs / newsgen Goto Github PK

Add script that will take processed dataset of headlines.json and captions.json and creates a train/validation/test split.

Update Readme with instruction to run Newsgen.

Right now TechCrunch is the only target news source in the crawl-input.json. We should add additional news sources to the list for a diverse news data.

Currently we are calculating warmup and total training steps based on data loader length. However, we need to consider the scenario of multi-device training in which case total length will be divided by number of devices.

Currently data loading is very slow. Speed can be improved by splitting the data files into train/val/test. Update encode_data.py. Also, remove raw text (headlines, captions, image paths, etc. to reduce size). Save using pytorch.save to compress it.

requirements.txt is empty. Add required dependencies to it.

Add generate method using greedy search or beam search.

Since current VQGAN model is using NN module and not Lightning module, the training loop should handle calls such as optimizer.step, loss.backward(), optimizer.zero_grad(), etc.

Using PyTorch Lightning, add an easy to use training/fine-tuning script.

create_subset.py takes a fixed subset size argument. Fix it to take subset size as percentage.

Use AdamW from transformers instead of Adam. Use linear schedule with warmup, transformers.get_linear_schedule_with_warmup.

Currently we are logging images every 1000 mini batches during validation and testing. However, we will not have enough images logged in case we have a small dataset or higher number of devices. Either we should change it to 100 mini-batches or take this value as a command line argument.

Original VQGAN implementation used PyTorch Lightning module but we converted to PyTorch NN module. It is much easier to do distributed training with PyTorch Lightning, so let's convert it to Lightning module.

Crawler needs to normalize urls. For example, https://www.example.com and https://www.example.com/ are the same and crawler shouldn't treat them as separate.

Using VQGAN encoder, convert all images in the given dataset to image tokens. These image tokens along with tokens from BART encoder (encoded captions/news headline) are fed to BART decoder to train it to generate image tokens given encoded captions/headlines.

We can do this as part of data transform step, however, doing it beforehand will speed up the training.