gemma

环境要求

• Python 环境（本项目为 Python 包，README 未指定具体版本）
• 已按 JAX 官网说明安装 CPU / GPU / TPU 版本的 JAX
• 可访问 PyPI 以执行 pip install gemma
• 可加载官方 Gemma checkpoint（示例通过 gm.ckpts 加载权重）

安装与启动步骤

1. 1安装 JAX

   先根据自身硬件（CPU、GPU 或 TPU）按 JAX 官网安装说明装好 JAX，Gemma 基于 JAX 运行，README 未给出具体命令。

2. 2安装 gemma 包

   在已装好 JAX 的环境中从 PyPI 安装官方 gemma 包，这是仓库对应的 Python 库。

   pip install gemma

   复制

3. 3构建模型并载权重

   导入 gemma 库，实例化 Gemma 4 模型，并按官方 checkpoint 路径加载预训练参数。

   from gemma import gm
   
   model = gm.nn.Gemma4_E4B()
   params = gm.ckpts.load_params(gm.ckpts.CheckpointPath.GEMMA4_E4B_IT)

   复制

4. 4多轮对话推理

   用 ChatSampler 开启多轮对话；示例为多模态，需自行准备 image1、image2 变量后再传入 images。

   sampler = gm.text.ChatSampler(
       model=model,
       params=params,
       multi_turn=True,
   )
   
   out0 = sampler.chat(prompt, images=[image1, image2])
   out1 = sampler.chat('What about the other image ?')

   复制

如何确认成功

把上面的示例代码保存为脚本运行，无导入或权重加载报错，并能打印出 sampler.chat 返回的模型回复，即表示安装与推理成功。

常见问题

注意事项

• README 只给出「先装 JAX、再 pip install gemma」两步安装说明，其余细节需查官方文档。
• 示例是图文多模态多轮对话，prompt 中使用 <|image|> 占位，并需提前准备好 image1、image2。
• Gemma 基于 Gemini 研究，仓库定位偏底层研究，代码清晰可读，也提供微调用例入口。