import argparse
import math
import os
import yaml
from tqdm import tqdm

from torch.utils.data import DataLoader
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    default_data_collator,
