RuntimeError: Error(s) in loading state_dict for MistralForCausalLM:
size mismatch for lm_head.base_layer.weight: copying a param with shape torch.Size([32002, 4096]) from checkpoint, the shape in current model is torch.Size([32000, 4096]).
size mismatch for lm_head.lora_B.default.weight: copying a param with shape torch.Size([32002, 16]) from checkpoint, the shape in current model is torch.Size([32000, 16]).
Iam gettting size mismatch error while i try to load my PEFT fine tuned model which has extra special tokens.
My code for loading the checkpoint is :
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
# Define the checkpoint directory
checkpoint_dir = "./mistral-json-finetune-Jul01/checkpoint-300" # Replace with your actual checkpoint ID
# Load the tokenizer from the checkpoint directory
eval_tokenizer = AutoTokenizer.from_pretrained(checkpoint_dir, trust_remote_code=True)
# Load the fine-tuned model from the checkpoint directory
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
ft_model = AutoModelForCausalLM.from_pretrained(
checkpoint_dir,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
)
# Resize the model's embeddings to match the tokenizer's vocabulary size
ft_model.resize_token_embeddings(len(eval_tokenizer))
# Ensure the model and tokenizer are synchronized
print(f"Tokenizer vocab size: {len(eval_tokenizer)}")
print(f"Model vocab size: {ft_model.config.vocab_size}")
# Example query
query = ''
while True:
query = input("Enter the query ('exit' to quit): ")
if query.lower() == 'exit':
break
# Construct the evaluation prompt
eval_prompt = f"""<s>[INST] You are an assistant which can extract key value in json format from a query text. Please create a json block for this input query text. Please do not add any explanation and show only json.
{query} [/INST] </s>
"""
# Tokenize and prepare input for the model
model_input = eval_tokenizer(eval_prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
# Generate response using fine-tuned model
ft_model.eval()
with torch.no_grad():
output_ids = ft_model.generate(**model_input, max_new_tokens=114, repetition_penalty=1.15)
response = eval_tokenizer.decode(output_ids[0], skip_special_tokens=True)[len(eval_prompt):].strip()
if len(response):
print('\n', response.strip(), '\n')
else:
print("\n{}\n")
print('Completed Execution')
How to resolve this.
Generating extra content repeted content after response till max_new_tokens is reached