[Python] Could not generate answer after finetuning bloomz-560m with LoRA

I have problems while generating the conversation after finetuning bloomz-560m with LoRA.

It seems okay while training, but it give me "IndexError: tuple index out of range" when generating answers. Also, the model could not work anymore no matter what input data I give.

Here is my code:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from peft import LoraConfig, get_peft_model, TaskType
from torch.utils.data import Dataset
from transformers import Trainer, TrainingArguments

MODEL_NAME = "bigscience/bloomz-560m"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
lora_config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
inference_mode=False,
r=8,
lora_alpha=16,
lora_dropout=0.1
)

model = get_peft_model(model, lora_config)
class CustomDataset(Dataset):
def __init__(self, input_texts):
self.input_texts = input_texts

def __len__(self):
return len(self.input_texts)

def __getitem__(self, idx):
input_text = self.input_texts[idx]
# Tokenization
encoding = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
# Include labels
encoding["labels"] = encoding["input_ids"]
return {key: val.squeeze() for key, val in encoding.items()}

input_texts = [
"Convesation 1:\nQ: 請問什麼是性騷擾？\nA: 性騷擾是指任何不受歡迎的性相關行為...",
"Convesation 2:\nQ: 如何避免性行為帶來的風險？\nA: 為了避免性行為帶來的健康風險...",
"Convesation 3:\nQ: 為什麼性教育對青少年很重要？\nA: 性教育對青少年來說很重要..."
]
dataset = CustomDataset(input_texts)
training_args = TrainingArguments(
output_dir="./lora_bloomz_finetuned",
per_device_train_batch_size=1,
num_train_epochs=3,
learning_rate=2e-5,
logging_steps=10,
save_steps=50,
save_total_limit=2,
report_to="none"
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset
)

trainer.train()
test_prompt = "Q: 什麼是避孕套的作用？\nA:"
inputs = tokenizer(test_prompt, return_tensors="pt").to(model.device)
output = model.generate(**inputs, max_length=100)
print(tokenizer.decode(output[0], skip_special_tokens=True))


Error message:

IndexError Traceback (most recent call last)
Cell In[10], line 3
1 test_prompt = "Q: 什麼是避孕套的作用？\nA:"
2 inputs = tokenizer(test_prompt, return_tensors="pt").to(model.device)
----> 3 output = model.generate(**inputs, max_length=100)
4 print(tokenizer.decode(output[0], skip_special_tokens=True))

File ~/Desktop/sex-edu-chatbot/.venv/lib/python3.10/site-packages/peft/peft_model.py:1704, in PeftModelForCausalLM.generate(self, *args, **kwargs)
1702 with self._enable_peft_forward_hooks(*args, **kwargs):
1703 kwargs = {k: v for k, v in kwargs.items() if k not in self.special_peft_forward_args}
-> 1704 outputs = self.base_model.generate(*args, **kwargs)
1705 else:
1706 outputs = self.base_model.generate(**kwargs)

File ~/Desktop/sex-edu-chatbot/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py:115, in context_decorator.<locals>.decorate_context(*args, **kwargs)
112 @functools.wraps(func)
113 def decorate_context(*args, **kwargs):
114 with ctx_factory():
--> 115 return func(*args, **kwargs)

File ~/Desktop/sex-edu-chatbot/.venv/lib/python3.10/site-packages/transformers/generation/utils.py:1829, in GenerationMixin.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
1826 batch_size = inputs_tensor.shape[0]
1828 device = inputs_tensor.device
-> 1829 self._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=device)
...
326 else:
327 # Note: don't use named arguments in `torch.isin`, see https://github.com/pytorch/pytorch/issues/126045
328 return torch.isin(elements, test_elements)

IndexError: tuple index out of range```


I tried using different input data, but it doesn't work.
And the pretrained model originally could translate conversation, but it could not do it after finetuning.

Continue reading...