model.py

import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

class BathSalt1DaedalusPhi3Model(AutoModelForSeq2SeqLM):
    def __init__(self, config):
        super().__init__(config)
        self.config = config

    def forward(self, input_ids, attention_mask, labels):
        outputs = self.model(input_ids, attention_mask=attention_mask, labels=labels)
        return outputs