Custom datasets support for chat models

[RanchiZhao]

for chat models, calibration datasets' input_ids&attn_masks should be passed in

[zengchao0424]

If you want to customize your own calibration dataset, you can add support for your dataset in the get_loader function inside ABQ-LLM/algorithm/datautils.py.For example, you can implement the get_chat function to load your own calibration dataset.

def get_loaders(
    name, nsamples=128, seed=0, seqlen=2048, model='',
):
    if 'wikitext2' in name:
        return get_wikitext2(nsamples, seed, seqlen, model)
    if 'chat_data' in name:
        return get_chat(nsamples, seed, seqlen, model)
    if 'ptb' in name:
        if 'new' in name:
            return get_ptb_new(nsamples, seed, seqlen, model)  
        return get_ptb(nsamples, seed, seqlen, model)
    if 'c4' in name:
        if 'new' in name:
            return get_c4_new(nsamples, seed, seqlen, model)  
        return get_c4(nsamples, seed, seqlen, model)
    if 'mix' in name:
        wiki_train,wiki_val=get_wikitext2(nsamples//3, seed, seqlen, model)
        ptb_train,ptb_val=get_ptb(nsamples//3, seed, seqlen, model)
        c4_train,c4_val=get_c4(nsamples//3, seed, seqlen, model)
        train=wiki_train+ptb_train+c4_train
        val=None
        return train,val