CACHE_DIR = ''          # Path to the cache directory
MODEL_DIR = ''          # Path to the model directory
INGREDIENTS_PATH = ""   # Path to the ingredients file (If exists)
PTM_PATH = ""           # Path to the pre-trained model

config = {
    'dataset': [
        {
            'name': 'snli',
            'batch_size': 64,
	    'mask_class': None,
        },

        {
            'name': 'mnli',
            'val_fraction': 0.2,
            'batch_size': 64,
	    'mask_class': None,
        },
        {
            'name': 'sick',
            'batch_size': 64,
	    'mask_class': None,
        },
        {
            'name': 'qnli',
            'val_fraction': 0.2,
            'batch_size': 64,
	    'mask_class': 1,
        },
        {
            'name': 'rte',
            'val_fraction': 0.5,
            'batch_size': 64,
	    'mask_class': 1,
        },
        {
            'name': 'scitail',
            'batch_size': 64,
	    'mask_class': 2,
        },
    
    ],
    'model': {
        'name' : 'meta-llama/Meta-Llama-3-8B',
        'ptm_path': PTM_PATH,
        'cachedir': './data/models--meta-llama--Meta-Llama-3-8B',
        'bases': [
            # Path to model checkpoints stored locally
            # f'{MODEL_DIR}/llama/selected/snli.pt',
            # f'{MODEL_DIR}/llama/selected/mnli.pt',
            # f'{MODEL_DIR}/llama/selected/sick.pt',
            # f'{MODEL_DIR}/llama/selected/qnli.pt',
            # f'{MODEL_DIR}/llama/selected/rte.pt',
            # f'{MODEL_DIR}/llama/selected/scitail.pt',

            #HF models IDs
            'hoffman-lab/KnOTS-Llama3_8B_lora_R16_snli',
            'hoffman-lab/KnOTS-Llama3_8B_lora_R16_mnli',
            'hoffman-lab/KnOTS-Llama3_8B_lora_R16_sick',
            'hoffman-lab/KnOTS-Llama3_8B_lora_R16_qnli',
            'hoffman-lab/KnOTS-Llama3_8B_lora_R16_rte',
            'hoffman-lab/KnOTS-Llama3_8B_lora_R16_scitail',
        ],
        'ft_config': {
            'type': 'lora',
            'subtype': 'peft',
            },


        'peft_config': {
            'task_type' : "SEQ_CLS",
            'inference_mode' : False,
            'r': 16,
            'lora_alpha' : 16,
            'lora_dropout' : 0.1,
            'target_modules' : ["q_proj", "k_proj", "v_proj", "o_proj"]
        },
    },
    'task_merge_config': {
        'representation': 'com-vector',
        'scaling_coeffs': 1.,
        'alpha_regmean': 0.6,
        'merge_method': 'regmean',
        'merge_split': "train",
        'gram_dtype': "torch.float32",
        'inv_gram_dtype': "torch.float32",
        'grams_data_distribution': "class-wise",  # or random
        'grams_data_measure': "quantity",   # or percentage
        'grams_data_quantity': 256,
        'grams_data_perc': 0.02,
        'weigh_tasks_by_grams': True,
        'cache_data': True
    },
    'eval_type': 'logits',
}

