fix: Changes in function process_dataargs to support the current implementation

[Abhishek-TAMU]

Description of the change

1- Adding unit test case for testing function process_dataargs with current way of tuning.
2- Changes to make the case of dataset with input, output key work with current way of tuning and make test_process_dataargs_pretokenized handler work.

Related issue number

https://github.ibm.com/ai-foundation/watson-fm-stack-tracker/issues/1428

How to verify the PR

Run current way of tuning for multiple cases:

• Case with using dataset with text and label as a single sequence, using this dataset

Command used:

python tuning/sft_trainer.py  \
--model_name_or_path Maykeye/TinyLLama-v0  \
--training_data_path tests/data/twitter_complaints_small.jsonl  \
--output_dir outputs/full-tuning  \
--num_train_epochs 5  \
--per_device_train_batch_size 2  \
--gradient_accumulation_steps 1  \
--learning_rate 1e-5  \
--response_template "\n### Label:"  \
--dataset_text_field "output" \
--use_flash_attn false \
--torch_dtype "float32"

• Case with pre-tokenized dataset using this dataset.

Command used:

python tuning/sft_trainer.py  \
--model_name_or_path Maykeye/TinyLLama-v0  \
--training_data_path tests/data/twitter_complaints_tokenized_with_maykeye_tinyllama_v0.jsonl \
--output_dir outputs/full-tuning  \
--num_train_epochs 5  \
--per_device_train_batch_size 2  \
--gradient_accumulation_steps 1  \
--learning_rate 1e-5  \
--use_flash_attn false \
--torch_dtype "float32"

• Case with input, output key in dataset using this dataset

Command used:

python tuning/sft_trainer.py  \
--model_name_or_path Maykeye/TinyLLama-v0  \
--training_data_path tests/data/twitter_complaints_input_output.jsonl  \
--output_dir outputs/full-tuning  \
--num_train_epochs 5  \
--per_device_train_batch_size 2  \
--gradient_accumulation_steps 1  \
--learning_rate 1e-5  \
--use_flash_attn false \
--torch_dtype "float32"

Was the PR tested

• I have added >=1 unit test(s) for every new method I have added.
• I have ensured all unit tests pass

[github-actions]

Thanks for making a pull request! 😃
One of the maintainers will review and advise on the next steps.

[Abhishek-TAMU]

Reference of discussion for this change: foundation-model-stack#381 (comment)

[dushyantbehl]

Thanks @Abhishek-TAMU

[Abhishek-TAMU]

In current implementation, dataset just have these columns input_ids, labels", attention_mask, hence every other columns in the dataset needed to be removed.