如何用LangChain为Llama2-7B编写自定义提示模板实现地址分类并修正格式

核心问题排查 #

你的代码存在几个关键问题导致模型输出不符合预期：

• 示例模板变量错误：example_formatter_template中City : {Address}是错误的变量引用，应该改为City : {City}，这会让模型学习到错误的字段映射逻辑。
• Prompt指令模糊：未明确强制模型返回指定格式的JSON，也未强调禁止添加额外信息。
• 缺少格式约束：未使用LangChain的输出解析器来确保输出结构合规，7B参数的小模型对格式的遵循度较弱。

分步解决方案 #

1. 修复示例模板的变量错误 #

修正example_formatter_template中的City字段变量，确保示例映射正确：

example_formatter_template = """
input: {input}
Address: {Address}
City: {City}
State: {State}
Country: {Country}
"""

2. 强化Prompt指令，明确格式要求 #

重新设计FewShotPromptTemplate的前缀和后缀，明确要求返回严格的JSON格式，禁止额外内容，缺失字段返回"未找到"：

prefix = """
你需要从输入的地址字符串中提取详细信息，严格按照以下规则处理：
1. 必须返回JSON格式，键为input、Address、City、State、Country
2. Address提取除城市、州、国家之外的具体地址部分
3. 若某字段不存在，返回"未找到"，禁止添加字符串中没有的信息
4. 不要输出任何额外的解释性文字，只返回JSON

以下是示例：
"""

suffix = """
现在处理输入：
input: {input}
请返回符合要求的JSON：
"""

few_shot_prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    prefix=prefix,
    suffix=suffix,
    input_variables=["input"],
    example_separator="\n---\n"
)

3. 添加JSON输出解析器强制格式 #

使用JSONOutputParser确保模型输出符合JSON结构，即使模型输出有微小偏差也能解析：

from langchain.output_parsers import JSONOutputParser
from pydantic import BaseModel, Field

# 定义输出结构的Pydantic模型
class AddressSchema(BaseModel):
    input: str = Field(description="输入的原始地址字符串")
    Address: str = Field(description="提取的具体地址部分")
    City: str = Field(description="提取的城市，不存在则返回'未找到'")
    State: str = Field(description="提取的州/省，不存在则返回'未找到'")
    Country: str = Field(description="提取的国家，不存在则返回'未找到'")

output_parser = JSONOutputParser(pydantic_object=AddressSchema)

# 将解析器的格式说明整合到prompt中
few_shot_prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    prefix=f"{prefix}\n{output_parser.get_format_instructions()}",
    suffix=suffix,
    input_variables=["input"],
    example_separator="\n---\n"
)

4. 调整LLM参数提升确定性 #

对于Llama2 7B模型，调整参数增强输出的确定性，减少随机性：

from langchain.llms import HuggingFacePipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=200,
    temperature=0.0,  # 设为0让输出更确定，避免随机内容
    top_p=1.0,
    repetition_penalty=1.1,
    return_full_text=False  # 只返回生成的内容，不包含输入prompt
)

llm = HuggingFacePipeline(pipeline=pipe)

5. 完整运行代码 #

整合所有修改后的完整代码：

from langchain.prompts import PromptTemplate, FewShotPromptTemplate
from langchain.chains import LLMChain
from langchain.output_parsers import JSONOutputParser
from pydantic import BaseModel, Field
from langchain.llms import HuggingFacePipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# 定义示例（优化第一个示例的Address字段，确保拆分逻辑准确）
examples = [
    {"input": "Plot No. 7, Sector 22, Noida, Uttar Pradesh, 201301, India",
     "Address": "Plot No. 7, Sector 22",
     "City" : "Noida",
     "State" : "Uttar Pradesh",
     "Country" : "India"},

    {"input": "Banjara Hills, Telangana, 500034, India",
     "Address": "Banjara Hills",
     "City" : "未找到",
     "State" : "Telangana",
     "Country" : "India"},
]

# 修复示例模板
example_formatter_template = """
input: {input}
Address: {Address}
City: {City}
State: {State}
Country: {Country}
"""
example_prompt = PromptTemplate(
    input_variables=["input", "Address", "City", "State", "Country"],
    template=example_formatter_template
)

# 定义输出结构
class AddressSchema(BaseModel):
    input: str = Field(description="输入的原始地址字符串")
    Address: str = Field(description="提取的具体地址部分")
    City: str = Field(description="提取的城市，不存在则返回'未找到'")
    State: str = Field(description="提取的州/省，不存在则返回'未找到'")
    Country: str = Field(description="提取的国家，不存在则返回'未找到'")

output_parser = JSONOutputParser(pydantic_object=AddressSchema)

# 构建Prompt
prefix = """
你需要从输入的地址字符串中提取详细信息，严格按照以下规则处理：
1. 必须返回JSON格式，键为input、Address、City、State、Country
2. Address提取除城市、州、国家之外的具体地址部分
3. 若某字段不存在，返回"未找到"，禁止添加字符串中没有的信息
4. 不要输出任何额外的解释性文字，只返回JSON

以下是示例：
"""

suffix = """
现在处理输入：
input: {input}
请返回符合要求的JSON：
"""

few_shot_prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    prefix=f"{prefix}\n{output_parser.get_format_instructions()}",
    suffix=suffix,
    input_variables=["input"],
    example_separator="\n---\n"
)

# 初始化LLM
model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=200,
    temperature=0.0,
    top_p=1.0,
    repetition_penalty=1.1,
    return_full_text=False
)

llm = HuggingFacePipeline(pipeline=pipe)

# 创建链并运行
chain = LLMChain(llm=llm, prompt=few_shot_prompt, verbose=True, output_parser=output_parser)
result = chain.run("B-12, Gandhi Colony, Bhopal, Madhya Pradesh, 462016, India")
print(result)

额外注意事项 #

• 优化了第一个示例的Address字段（去掉Noida），让模型学习更准确的拆分逻辑。
• 若模型仍有格式问题，可以尝试增加更多示例，尤其是包含缺失字段的场景。
• 确保使用Llama2的chat版本模型（如Llama-2-7b-chat-hf），该版本对指令的遵循度更高。

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