Key error on Week1

[treblenalto]

Hi @graviraja,
I was following your tutorial on wandb logging and found a potential error in the training code when visualizing poorly performed data with wandb table.

MLOps-Basics/week_1_wandb_logging/train.py

Lines 19 to 21 in 13060e0

	def on_validation_end(self, trainer, pl_module):
	val_batch = next(iter(self.datamodule.val_dataloader()))
	sentences = val_batch["sentence"]

When running, this results in KeyError : "sentence", referring to line 21 sentences = val_batch["sentence"].

I think this is because "sentence" is not part of the columns setup in val_data of data.py. Please correct me if I'm wrong. Thanks :)

MLOps-Basics/week_1_wandb_logging/data.py

Lines 29 to 42 in 13060e0

	def setup(self, stage=None):
	# we set up only relevant datasets when stage is specified
	if stage == "fit" or stage is None:
	self.train_data = self.train_data.map(self.tokenize_data, batched=True)
	self.train_data.set_format(
	type="torch", columns=["input_ids", "attention_mask", "label"]
	)
	self.val_data = self.val_data.map(self.tokenize_data, batched=True)
	self.val_data.set_format(
	type="torch",
	columns=["input_ids", "attention_mask", "label"],
	output_all_columns=True,
	)

[PTson2207]

Hi @Taehee-K ,
I also got the same error as you. Did you fix that error?
Can you show me the results?

[agamemnonc]

The solution to this is just to also include sentence in your data batching:

self.train_data.set_format( 
             type="torch", columns=["input_ids", "attention_mask", "label", "sentence"] 

and the same obviously for self.val_data.set_format() and self.test_data.set_format().