Gemma LLM Library

The Sudoku preprocessing directory contains the following components:

Scripts:

• prepare_sudoku_dataset.sh: Configures Kaggle credentials and fetches the dataset.
• convert_sudoku.py: Preprocesses CSV raw datasets into Bagz files.
• sudoku_data.py: Handles dataset and transform configuration using Kauldron pipelines.

Generated Files:

• sudoku_train.bagz: Converted training records containing source puzzles and reference solutions.
• sudoku_eval.bagz: Converted evaluation records.

The repository contains non-official research projects located in the gemma/research/ directory. One such project is:

• t5gemma: An encoder/decoder Gemma architecture based on Gemma 2.

To validate core architectures, custom layers, losses, and utilities, run the following command:

pytest gemma/diffusion/hackable_diffusion_adapter/hd/...

To use the Hackable Diffusion Adapter, you must install the gemma package and the jax[cuda13] dependencies. It is recommended to use Python 3.12 and CUDA 13.

Warning: Do not mix CUDA 13 with CUDA 12 packages (such as jax-cuda12-plugin or nvidia-nccl-cu12) as this can cause NCCL errors.

# Install from PyPI
pip install gemma

# Or install from source
git clone https://github.com/google-deepmind/gemma.git
cd gemma
pip install .

# Install required JAX dependencies
pip install -U jax[cuda13]

Use the pubmedqa_eval.py module to evaluate model performance on the PubMedQA dataset. This includes extracting structured answers and calculating accuracy or BLEU scores.

• extract_pubmedqa_answer: Extracts the structured answer (yes/no/maybe) from a model response by searching for the marker: "The answer is: {yes|no|maybe}".
• PubMedQAAccuracy: Computes the accuracy of extracted answers against ground truth.
• BLEUScore: Computes smoothed sentence-level BLEU scores between generated and ground-truth text using sacrebleu.

Evaluation is an offline process run after training. It loads checkpoints and performs autoregressive (AR) sampling to report metrics like accuracy (Sudoku) or BLEU (PubMedQA).

Run the evaluation command from the parent directory of the gemma directory.

env XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" \
    XLA_PYTHON_CLIENT_PREALLOCATE="false" \
    TF_FORCE_GPU_ALLOW_GROWTH="true" \
    python3 -m gemma.diffusion.hackable_diffusion_adapter.eval_main \
    --cfg=gemma/diffusion/hackable_diffusion_adapter/configs/sft_sudoku.py \
    --task=sudoku \
    --step=1000 \
    --eval_names=sample_ar_steps64 \
    --cfg.workdir=$(pwd)/xp_dir_sudoku_lora \
    --cfg.eval_ds.batch_size=2 \
    --cfg.aux.eval_num_batches=2 \
    --cfg.aux.num_canvases=2

You can replace existing modules with their LoRA or quantized versions using a ModuleInterceptor. This allows you to transform a model's architecture dynamically during a forward pass.

# Replace dense layers with LoRA
def _replace_dense_by_lora(module: nn.Module) -> nn.Module:
  if isinstance(module, nn.Dense):
    return peft.LoRADense(rank=3, wrapped=module)
  else:
    return module

with ModuleInterceptor(_replace_dense_by_lora):
  y = model(x)

To download Gemma models manually from Kaggle Hub:

1. Navigate to the desired model page (e.g., Gemma 4, Gemma 3, etc.).
2. Select one of the Flax model variations.
3. Click the "Download" button to obtain the model archive.
4. Extract the archive.

The extracted archive contains both the model weights and the tokenizer (e.g., a directory for weights and a tokenizer.model file).