Updated LIT + Gemma Colab for LIT v1.2.

[RyanMullins]

Description of the change

Updates to the LIT + Gemma Colab to use the new functionality introduced in LIT v1.2.

Motivation

These changes address user feedback collected from Gemma DevRel and reported to the LIT team.

Type of change

Choose one: Documentation

Checklist

• I have performed a self-review of my code.
• I have added detailed comments to my code where applicable.
• I have verified that my change does not break existing code.
• My PR is based on the latest changes of the main branch (if unsure, please run git pull --rebase upstream main).
• I am familiar with the Google Style Guide for the language I have coded in.
• I have read through the Contributing Guide and signed the Contributor License Agreement.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[MarkDaoust]

@gustheman