Update model.py

src/models/model.py

@@ -79,3 +79,8 @@ class DataModule(Dataset):
             text=data_row["text"],
             keywords_input_ids=input_encoding["input_ids"].flatten(),
             keywords_attention_mask=input_encoding["attention_mask"].flatten(),
+            labels_attention_mask=output_encoding["attention_mask"].flatten(),
+        )
+
+
+class PLDataModule(LightningDataModule):