From c76b63772c6cb3edba8c5dbbb45a9006e4dee73f Mon Sep 17 00:00:00 2001 From: psaegert Date: Sun, 9 Feb 2025 10:14:15 +0100 Subject: [PATCH] Fix labels missing indexing --- src/flash_ansr/compat/evaluation_nesymres.py | 2 +- src/flash_ansr/compat/evaluation_pysr.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/flash_ansr/compat/evaluation_nesymres.py b/src/flash_ansr/compat/evaluation_nesymres.py index 961214e..1db90b9 100644 --- a/src/flash_ansr/compat/evaluation_nesymres.py +++ b/src/flash_ansr/compat/evaluation_nesymres.py @@ -112,7 +112,7 @@ def evaluate( results_dict['n_support'].append([batch['x_tensors'].shape[1] // 2] * batch['x_tensors'].shape[0]) # Create the labels for the next token prediction task (i.e. shift the batch['input_ids'] by one position to the right) - labels = batch['labels'].clone() + labels = batch['labels'][0].clone() labels_decoded = expression_space.tokenizer.decode(labels.tolist(), special_tokens='') # TODO: For different datasets, sort unused dimensions to the end diff --git a/src/flash_ansr/compat/evaluation_pysr.py b/src/flash_ansr/compat/evaluation_pysr.py index dbf1d94..2c52fc9 100644 --- a/src/flash_ansr/compat/evaluation_pysr.py +++ b/src/flash_ansr/compat/evaluation_pysr.py @@ -132,7 +132,7 @@ def evaluate( results_dict['n_support'].append([x_tensor.shape[1] // 2] * x_tensor.shape[0]) # Create the labels for the next token prediction task (i.e. shift the input_ids by one position to the right) - labels = batch['labels'].clone() + labels = batch['labels'][0].clone() labels_decoded = expression_space.tokenizer.decode(batch['labels'].tolist(), special_tokens='') # TODO: For different datasets, sort unused dimensions to the end