Unable to run downstream_fine_tune.ipynb

[moa4020]

Hello,

Thank you very much for your work! I'm excited to run your model on our spatial data and finetune it for niche regression. So, I first used the get_embeddings notebook to embed one of my datasets. And now, I'm using the fine tuning notebook to train a model that can predict niche labels.

But I've been facing three issues:

1. I've read the paper, but I'm unsure if cell-typing is a feature in nicheformer. Because the setting for supervised_task in config seems to take different inputs like 'classification', 'regression', 'niche_regression', 'density_regression', 'niche_classification', 'niche_binary_classification', 'niche_multiclass_classification'. Are these all only associated with niche/spatial related predictions? PS: I got these looking up supervised_task in your codebase. Please correct me if my assumptions are wrong and help clarify what each option does.
2. I'm not sure what the baseline=False setting is supposed to do in NicheformerFineTune(). So I commented it out because I was getting an error saying it was an invalid argument.
3. And finally, when I run the notebook to finetune a model for niche_regression, I cannot get past the sanity check and end up with keyError: 'modality' at the end of the error stack.

For reproducibility, I have attached some sample data and the html rendering of my notebook. Any advice is highly appreciated!

Thank you,
Mar

03_nicheformer_downstream_fine_tune.html

sample_data.h5ad.zip