Fix tests for custom dataset, grammar, batching, chat_completion

src/tests/datasets/test_custom_dataset.py

src/tests/datasets/test_grammar_datasets.py

src/tests/test_batching.py

src/tests/test_chat_completion.py