Add prediction (table) artifacts to Weights & Biases logger

[Glavin001]

⚠️ Please check that this feature request hasn't been suggested before.

• [X] I searched previous Ideas in Discussions didn't find any similar feature requests.
• [X] I searched previous Issues didn't find any similar feature requests.

🔖 Feature description

One of my favourite features from LLM Studio is the validation prediction insights: https://h2oai.github.io/h2o-llmstudio/guide/experiments/view-an-experiment#experiment-tabs

Validation prediction insights : This tab displays model predictions for random, best, and worst validation samples. This tab becomes available after the first validation run and allows you to evaluate how well your model generalizes to new data.

Since Axolotl is headless (no UI) this can instead be implemented with WandB logging.

Examples:

• https://wandb.ai/ayush-thakur/plant-pathology-21/artifacts/val_prediction/val_predictions/v1/files/val_prediction.table.json
• 

✔️ Solution

See https://wandb.ai/stacey/mnist-viz/reports/Visualize-Predictions-over-Time--Vmlldzo1OTQxMTk

❓ Alternatives

No response

📝 Additional Context

I'd be interested in contributing this, if Axolotl team is interested and I can figure it out 😅

Acknowledgements

• [X] My issue title is concise, descriptive, and in title casing.
• [X] I have searched the existing issues to make sure this feature has not been requested yet.
• [X] I have provided enough information for the maintainers to understand and evaluate this request.

[NanoCode012]

A callback could be added for this feature.

~I wasn’t sure wandb support saving text results.~

Edit: Wandb table can save predictions.

[Glavin001]

@NanoCode012 : Could you give me some pointers on where this should be added to Axolotl? I'll try to find time in the next week when I'm training to add and test this new feature. Thanks!

[NanoCode012]

callbacks should be placed in utils/callbacks.py. Then you can add it to the Trainer at utils/trainer.py. You can see examples of callbacks and how it's added in the aforementioned files.

I think you can add a callback on_evaluate finished (?) if that's an option to also predict over a few eval samples and save the responses.