如何让google/gemma-2-2b在Hugging Face Transformers中严格遵循指令？

问题场景 #

调用google/gemma-2-2b模型时，模型未严格遵循指令，反而重复提示词并生成额外选项，无法输出预期的单一选项结果。执行代码如下：

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch

HUGGINGFACE_TOKEN = "<my-token>"

model_name = "google/gemma-2-2b"
tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_TOKEN)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16,
    token=HUGGINGFACE_TOKEN
)
text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer, token=HUGGINGFACE_TOKEN)

prompt = "What is the capital of France? Just select an option. Choose only one option from the following A) Paris B) London C) Delhi 4) Goa"
output = text_generator(prompt, max_new_tokens=100)
print(output)

预期输出：

A) Paris

实际输出：

[{'generated_text': 'What is the capital of France? Just select an option. Choose only one option from the following A) Paris B) London C) Delhi 4) Goa 5) New York ...'}]

解决方案 #

可以从提示词优化和生成参数调整两方面入手，让模型严格遵循指令输出：

一、提示词优化 #

• 明确输出约束：在提示词末尾添加清晰的输出要求，避免歧义，比如：

  请从以下选项中选择法国首都对应的唯一答案，直接输出选项内容，不要添加任何额外文字：A) Paris B) London C) Delhi D) Goa
  

• 统一选项格式：将编号统一为字母（把4)改为D)），防止模型混淆编号逻辑，减少额外生成选项的可能。
• 结构化指令：用更清晰的任务描述引导模型，比如：

  任务：选择法国首都的正确选项
  要求：仅输出选项编号和内容，无其他文字
  选项：
  A) Paris
  B) London
  C) Delhi
  D) Goa
  

二、生成参数调整 #

调整pipeline的生成参数，限制模型的发散性输出：

• temperature=0：完全消除随机性，让模型只选择概率最高的token，避免生成无关内容。
• do_sample=False：启用贪婪解码，直接选取概率最高的输出序列，确保结果确定性。
• max_new_tokens=5：设置较小的最大生成token数，防止模型生成冗长内容。
• return_full_text=False：仅返回模型新生成的文本，不包含原始提示词，方便直接获取结果。
• stop_sequence=["\n"]：指定停止生成的触发符，当模型输出换行时立即停止，避免多余内容。

修改后的完整代码 #

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch

HUGGINGFACE_TOKEN = "<my-token>"

model_name = "google/gemma-2-2b"
tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_TOKEN)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16,
    token=HUGGINGFACE_TOKEN
)
text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer, token=HUGGINGFACE_TOKEN)

# 优化后的提示词
prompt = "请从以下选项中选择法国首都对应的唯一答案，直接输出选项内容，不要添加任何额外文字：A) Paris B) London C) Delhi D) Goa"
output = text_generator(
    prompt,
    max_new_tokens=5,
    temperature=0,
    do_sample=False,
    return_full_text=False,
    stop_sequence=["\n"]
)
# 提取并清理输出结果
print(output[0]['generated_text'].strip())

执行上述代码后，模型将输出符合要求的A) Paris。

内容的提问来源于stack exchange，提问作者hanugm