Finetuning llama-2 for classification

used AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    token=token
    )

and

bnb_config = BitsAndBytesConfig(
    load_in_4bit = True,
    bnb_4bit_qunat_type = "nf4",
    bnb_4bit_compute_dtype = torch.float16,
)

for loading the model and finetuned this model by using LoRA

peft_config = LoraConfig(
        r=16,
        lora_alpha=64,
        lora_dropout=0.1,
        bias="none",
        task_type='SEQ_CLS',
)

and saved as “tuned_model”
while loading the model:

from transformers import pipeline
pipe = pipeline('text-classification',
                tuned_model,
                device_map="auto")
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
[<ipython-input-19-fd1b8ca97698>](https://localhost:8080/#) in <cell line: 2>()
      1 from transformers import pipeline
----> 2 pipe = pipeline('text-classification',tuned_model, device_map="auto")

5 frames
[/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py](https://localhost:8080/#) in <dictcomp>(.0)
   3809                 p: {"safetensors_file": f, "weight_name": p, "dtype": str_dtype}
   3810                 for p, f in weight_map.items()
-> 3811                 if param_device_map[p] == "disk"
   3812             }
   3813 

KeyError: 'lm_head.weight'

Can any one suggest me how to load this tuned_model ?

@zuhashaik were you able to solve this? I encountered the same issue but wasn’t able to find a solution

No, I stopped using pipeline, instead go ahead with the Automodelforsequenceclassification class it will recognise your adapter and load your tuned model and get the logits.
Hope this helps! Any further problem dm me zuhashaik12@gmail.com happy to help! :slight_smile:

1 Like