feat!: model.predict returns all the columns

[GarrettWu]

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:

• Make sure to open an issue as a bug/issue before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea
• Ensure the tests and linter pass
• Code coverage does not decrease (if any source code was changed)
• Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕

[ashleyxuu]

Do we want to make sure the generative-ai notebooks are edited accordingly?

[GarrettWu]

Do we want to make sure the generative-ai notebooks are edited accordingly?

Unless we add those into our tests, at some point it (the notebooks) will become unavailable. Before that, we can't cover those notebooks all the time.