Fix the ValueError while loading the Gemma model in logit_distillation.ipynb

[rajasekharporeddy]

Currently, the logit_distillation.ipynb notebook fails with the following ValueError:

AttributeError                            Traceback (most recent call last)
File /usr/local/lib/python3.12/site-packages/tunix/models/gemma/model.py:839, in Gemma.from_params(cls, params, version)
    838 try:
--> 839   config = getattr(ModelConfig, config_id)()
    840 except AttributeError as exc:

AttributeError: type object 'ModelConfig' has no attribute 'gemma_1_1_7b_it'

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
Cell In[5], line 29
     25   print(f"Finished processing {model_handle}.")
     28 # Load Teacher Model (Gemma 7B)
---> 29 load_and_save_model(
     30     "google/gemma/flax/1.1-7b-it", "1.1-7b-it", TEACHER_CKPT_DIR
     31 )
     33 # Load Student Model (Gemma 2B)
     34 load_and_save_model(
     35     "google/gemma/flax/1.1-2b-it", "1.1-2b-it", STUDENT_CKPT_DIR
     36 )

Cell In[5], line 13, in load_and_save_model(model_handle, version, ckpt_dir)
      9 with jax.default_device(jax.devices("cpu")[0]):
     10   params = params_lib.load_and_format_params(
     11       os.path.join(kaggle_ckpt_path, ckpt_version)
     12   )
---> 13   gemma = gemma_lib.Gemma.from_params(params, version=version)
     15 print(f"Saving checkpoint to {ckpt_dir}...")
     16 checkpointer = ocp.StandardCheckpointer()

File /usr/local/lib/python3.12/site-packages/tunix/models/gemma/model.py:841, in Gemma.from_params(cls, params, version)
    839   config = getattr(ModelConfig, config_id)()
    840 except AttributeError as exc:
--> 841   raise ValueError(f'Unsupported version: {version}') from exc
    843 return module_from_linen_variables(
    844     module_factory=lambda: cls(config, rngs=nnx.Rngs(params=0)),
    845     variables=params['transformer'],
    846     map_key_fn=_map_linen_var_names,
    847     assign_val_fn=_assign_linen_params_to_nnx_state,
    848 )

ValueError: Unsupported version: 1.1-7b-it

while loading the Gemma-1.1-7b-it model.

This PR fixes this error and with this change the notebook runs successfully.

[wang2yn84]

Thank you very much!