Как дообучить Llama 3 для задачи регрессии?

Я пытаюсь дообучить Llama-3-8B для задачи регрессии с использованием PEFT, чтобы модель выдавала число с плавающей запятой для “оценки” входных данных. Я попытался адаптировать модель, изменив количество выходных логитов на 1 и изменив функцию потерь на среднеквадратическую ошибку. Я отформатировал свои данные, предоставив инструкцию вместе с входным запросом, токенизировал их и сохранил значение с плавающей запятой в “labels”. Я получаю ошибку:

RuntimeError: The size of tensor a (8) must match the size of tensor b (512) at non-singleton dimension 1

где моя MAX_SEQ_LEN при токенизации входных инструкций и запросов установлена на максимальную длину 512, но как-то это совпадает с моим выходным значением с плавающей запятой? Все ресурсы в интернете для дообучения Llama 3 предназначены для текстовых ответов, а не для задач регрессии, поэтому я не уверен, как это сделать. Ниже приведён мой код:

from datasets import load_dataset, load_metric
from transformers import AutoTokenizer, Trainer, TrainingArguments, EarlyStoppingCallback, AutoModelForSequenceClassification, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model
from trl import SFTTrainer
import numpy as np
import json
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import torch

quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4"
)

MAX_SEQ_LEN = 512
device = "cuda" if torch.cuda.is_available() else "cpu"

dataset = load_dataset("json", data_files="datasets/aNLI/train_rationale_cleaned.jsonl", cache_dir="/data/etey916")["train"]
#dataset = load_dataset("json", data_files="datasets/aNLI/train_rationale_output.jsonl", cache_dir="/data/etey916")
model_name = "meta-llama/Meta-Llama-3-8B"
model = AutoModelForSequenceClassification.from_pretrained(model_name, quantization_config=quant_config, problem_type="regression", num_labels=1)
model.config.pad_token_id = model.config.eos_token_id
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

def format_dataset(sample):
    input_text = sample['prompt'].replace('_', sample['rationale'])
    instruction = ("<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n"
            "Приведите заявление, состоящее из начального контекста, гипотезы, вывода и результата, определите, насколько это правдоподобно, присвоив непрерывный балл от -1 до 1, где -1 означает чрезвычайно неправдоподобно, а 1 означает чрезвычайно правдоподобно.\n"
            f"{input_text}<|eot_id|>")
    output = float(sample['label'])
    tokenized_inputs = tokenizer(instruction, truncation=True, padding="max_length", max_length=MAX_SEQ_LEN)
    return {**tokenized_inputs, "labels": output}

def compute_metrics_for_regression(eval_pred):
    logits, labels = eval_pred
    labels = labels.reshape(-1, 1)
    mse = mean_squared_error(labels_logits)
    mae = mean_absolute_error(labels, logis)
    r2 = r2_score(labels, logits)
    sse = ((logits - labels).flatten()**2).tolist()

    accuracy = sum([1 for e in sse if e < 0.25]) / len(sse)

    return {"mse":mse, "mae":mae, "r2":r2, "accuracy":accuracy}

dataset = dataset.map(
    format_dataset,
    remove_columns=dataset.column_names
)

lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules="all-linear",
    bias = "none",
    task_type="CAUSAL_LM",
)

peft_model = get_peft_model(model, lora_config)

training_args = TrainingArguments(
    output_dir = "./results",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=5e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=10,
    weight_decay=0.01,
    logging_dir="./logs",
    logging_steps=10,
    save_steps=500,
    save_total_limit=2,
    #metric_for_best_model="eval_loss",
    load_best_model_at_end=True,
    greater_is_better=False
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    peft_config=lora_config,
    train_dataset=dataset,
    compute_metrics=compute_metrics_for_regression,
)

trainer.train()
trainer.save_model("./finetuned_model")
tokenizer.save_pretrained("./finetuned_model")