Update model.py

src/models/model.py

@@ -168,3 +168,20 @@ class LightningModel(LightningModule):
         output = self.model(
             input_ids,
             attention_mask=attention_mask,
+            labels=labels,
+            decoder_attention_mask=decoder_attention_mask,
+        )
+
+        return output.loss, output.logits
+
+    def training_step(self, batch, batch_size):
+        """ training step """
+        input_ids = batch["keywords_input_ids"]
+        attention_mask = batch["keywords_attention_mask"]
+        labels = batch["labels"]
+        labels_attention_mask = batch["labels_attention_mask"]
+
+        loss, outputs = self(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            decoder_attention_mask=labels_attention_mask,