采样#
演示如何加载 Gemma 模型并在其上运行推理的示例。
Gemma 库提供了 3 种提示模型的方法
gm.text.ChatSampler：最易于使用，只需与模型对话即可获得答案。开箱即用地支持多轮对话。gm.text.Sampler：级别较低，但提供更多控制。必须手动处理聊天状态以进行多轮对话。model.apply：直接调用模型，仅预测单个 token。
!pip install -q gemma
# Common imports
import os
import jax
import jax.numpy as jnp
# Gemma imports
from gemma import gm
默认情况下，Jax 不会完全利用 GPU 内存，但这可以被覆盖。请参阅 GPU 内存分配
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]="1.00"
加载模型和参数。这里我们加载模型的指令调优版本。
model = gm.nn.Gemma3_4B()
params = gm.ckpts.load_params(gm.ckpts.CheckpointPath.GEMMA3_4B_IT)
多轮对话#
与 Gemma 聊天的最简单方法是使用 gm.text.ChatSampler。它隐藏了对话缓存的样板代码，以及用于格式化对话的 <start_of_turn> / <end_of_turn> token。
在这里，我们在创建 gm.text.ChatSampler 时设置 multi_turn=True（默认情况下，ChatSampler 每次都会开始新的对话）。
在多轮模式下，您可以通过传递 chatbot.chat(..., multi_turn=False) 来擦除之前的对话状态。
sampler = gm.text.ChatSampler(
model=model,
params=params,
multi_turn=True,
)
turn0 = sampler.chat('Share one methapore linking "shadow" and "laughter".')
print(turn0)
Okay, here's a metaphor linking "shadow" and "laughter," aiming for a slightly evocative and layered feel:
**"Laughter is the fleeting shadow of joy, dancing across a face that’s often hidden in the long shadow of sorrow."**
---
**Here's a breakdown of why this works:**
* **"Shadow"** represents sadness, pain, or a past experience that lingers. It’s not necessarily a dark shadow, but a persistent presence.
* **"Laughter"** is presented as a brief, bright appearance – a momentary flash of happiness.
* **"Dancing across a face that’s often hidden"** emphasizes that the joy isn't constant, and the underlying sadness is still there, obscuring it.
---
Would you like me to:
* Try a different type of metaphor?
* Expand on this one with a short story snippet?
turn1 = sampler.chat('Expand it in a haiku.')
print(turn1)
Okay, here’s a haiku based on the metaphor:
Shadow stretches long,
Laughter’s brief, bright, dancing grace,
Joy hides in the dark.
---
Would you like me to try another haiku, or perhaps a different poetic form?
注意：默认情况下（multi_turn=False），对话状态每次都会重置，但您仍然可以通过传递 sampler.chat(..., multi_turn=True) 来继续之前的对话
默认情况下，使用贪婪解码。您可以传递自定义 sampling= 方法作为 kwargs
gm.text.Greedy()：（默认）贪婪解码gm.text.RandomSampling()：简单的随机采样和温度参数，以获得更多样性
采样提示#
为了获得更多控制，我们还提供了 gm.text.Sampler，它仍然执行高效采样（具有 kv-caching、提前停止等）。
提示采样器需要使用 <start_of_turn> / <end_of_turn> token 正确格式化提示（请参阅 分词器 上的自定义 token 部分文档）。
sampler = gm.text.Sampler(
model=model,
params=params,
)
prompt = """<start_of_turn>user
Give me a list of inspirational quotes.<end_of_turn>
<start_of_turn>model
"""
out = sampler.sample(prompt, max_new_tokens=1000)
print(out)
Okay, here's a list of inspirational quotes, categorized a little to give you a variety:
**On Perseverance & Resilience:**
* “The only way to do great work is to love what you do.” – Steve Jobs
* “Fall seven times, stand up eight.” – Japanese Proverb
* “The difference between ordinary and extraordinary is that little extra.” – Jimmy Johnson
* “Success is not final, failure is not fatal: It is the courage to continue that counts.” – Winston Churchill
* “Don’t watch the clock; do what it does. Keep going.” – Sam Levenson
* “When the going gets tough, the tough get going.” – Theodore Roosevelt
**On Self-Love & Confidence:**
* “You are enough.” – Brené Brown
* “Believe you can and you’re halfway there.” – Theodore Roosevelt
* “You must be the change you wish to see in the world.” – Mahatma Gandhi
* “The best is yet to come.” – Frank Sinatra
* “Be the energy you want to attract.” – Tony Gaskins
* “Don’t be defined by your past. Define your future.” – Unknown
**On Dreams & Goals:**
* “If you can dream it, you can do it.” – Walt Disney
* “The future belongs to those who believe in the beauty of their dreams.” – Eleanor Roosevelt
* “Shoot for the moon. Even if you miss, you’ll land among the stars.” – Les Brown
* “Start where you are. Use what you have. Do what you can.” – Arthur Ashe
* “Life begins at the end of your comfort zone.” – Unknown
**On Happiness & Perspective:**
* “Happiness is not something readymade. It comes from your own actions.” – Dalai Lama
* “It’s not the triumph that matters, it’s the effort.” – Winston Churchill
* “Don’t wait for the perfect moment, take the moment and make it perfect.” – Oscar Wilde
* “Be present. Be grateful. Be you.” – Unknown
* “The only way out is through.” – Robert Frost
**Short & Powerful:**
* “Be the change.” – Mahatma Gandhi
* “Just breathe.”
* “Keep going.”
* “You got this.”
* “Dream big.”
---
**Resources for More Quotes:**
* **BrainyQuote:** [https://www.brainyquote.com/](https://www.brainyquote.com/)
* **Goodreads:** [https://www.goodreads.com/quotes](https://www.goodreads.com/quotes)
* **Quote Garden:** [https://quotegarden.com/](https://quotegarden.com/)
To help me give you even more relevant quotes, could you tell me:
* **What kind of inspiration are you looking for?** (e.g., motivation for work, overcoming challenges, self-love, etc.)
* **Is there a particular theme or topic you'd like quotes about?**<end_of_turn>
直接使用模型#
这是一个直接调用模型预测单个 token 的示例。
模型输入期望编码后的 token。为此，我们首先需要使用分词器对提示进行编码。有关使用分词器的更多信息，请参阅我们的 分词器 文档。
tokenizer = gm.text.Gemma3Tokenizer()
注意：在编码提示时，不要忘记使用 add_bos=True 添加字符串开始 token。所有馈送到模型的提示都应该以此 token 开头。
prompt = tokenizer.encode('One word to describe Paris: \n\n', add_bos=True)
prompt = jnp.asarray(prompt)
然后我们可以调用模型，并获取预测的 logits。
# Run the model
out = model.apply(
{'params': params},
tokens=prompt,
return_last_only=True, # Only predict the last token
)
# Sample a token from the predicted logits
next_token = jax.random.categorical(
jax.random.key(1),
out.logits
)
tokenizer.decode(next_token)
'Romantic'
您还可以显示下一个 token 的概率。
tokenizer.plot_logits(out.logits)