使用 Hugging Face Transformers 和 PyTorch 实现信息抽取
在自然语言处理(NLP)领域,信息抽取是一种常见的任务,其目标是从文本中提取特定类型的结构化信息。本文将介绍如何使用 Hugging Face Transformers 和 PyTorch 实现基于大语言模型的信息抽取。我们将通过一个具体的例子,展示如何从文本中抽取商品的详细信息,包括产品名称、品牌、特点、价格和销量等。
1. 任务背景
信息抽取在许多应用场景中都非常有用,例如从新闻文章中提取关键信息、从产品描述中提取商品信息等。传统的信息抽取方法通常依赖于规则匹配或基于机器学习的分类器,但这些方法往往需要大量的标注数据,并且在面对复杂的文本结构时效果有限。近年来,随着大语言模型(如 GPT、Bert 等)的发展,基于这些模型的信息抽取方法逐渐成为研究热点。
2. 技术栈介绍
2.1 Hugging Face Transformers
Hugging Face Transformers 是一个开源的 Python 库,提供了大量预训练的语言模型,以及用于微调和部署这些模型的工具。它支持多种模型架构(如 Bert、GPT、T5 等),并且提供了丰富的 API,方便开发者快速实现各种 NLP 任务。
2.2 PyTorch
PyTorch 是一个流行的深度学习框架,以其动态计算图和易用性而闻名。它提供了强大的 GPU 支持,能够加速模型的训练和推理过程。在本文中,我们将使用 PyTorch 来加载和运行预训练的语言模型。
3. 实现步骤
3.1 定义任务和数据
我们的目标是从给定的文本中抽取商品的相关信息。具体来说,我们需要提取以下属性:
- 产品名称
- 品牌
- 特点
- 原价
- 促销价
- 销量
为了实现这一目标,我们定义了一个简单的数据结构来描述这些属性,并提供了一些示例数据供模型学习。
3.2 初始化模型和分词器
我们使用 Hugging Face Transformers 提供的 AutoTokenizer
和 AutoModelForCausalLM
来加载预训练的语言模型。在这个例子中,我们使用了 Qwen2.5-1.5B-Instruct
模型,这是一个经过指令微调的模型,能够更好地理解自然语言指令。
tokenizer = AutoTokenizer.from_pretrained(r"C:\Users\妄生\Qwen2.5-1.5B-Instruct", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(r"C:\Users\妄生\Qwen2.5-1.5B-Instruct", trust_remote_code=True)
3.3 构建提示模板
为了帮助模型更好地理解任务,我们设计了一个提示模板,明确告诉模型需要抽取哪些信息。模板中包含了需要抽取的实体类型和属性,以及如何格式化输出结果。
IE_PATTERN = "{}\n\n提取上述句子中{}的实体,并按照JSON格式输出,上述句子中不存在的信息用['原文中未提及']来表示,多个值之间用','分隔。"
3.4 准备示例数据
为了让模型更好地理解任务,我们提供了一些示例数据。这些示例数据展示了如何从文本中抽取信息,并以 JSON 格式输出。
ie_examples = {'商品': [{'content': '2024 年新款时尚运动鞋,品牌 JKL,舒适透气,多种颜色可选。原价 599 元,现在促销价 499 元。月销量 2000 双。','answers': {'产品': ['时尚运动鞋'],'品牌': ['JKL'],'特点': ['舒适透气、多种颜色可选'],'原价': ['599元'],'促销价': ['499元'],'月销量': ['2000双'],}}]
}
3.5 实现信息抽取函数
我们定义了一个 inference
函数,用于对输入的文本进行信息抽取。这个函数首先将输入文本和提示模板组合起来,然后使用模型生成输出结果。最后,我们将生成的结果解码并打印出来。
def inference(sentences: list, custom_settings: dict):for sentence in sentences:cls_res = "商品"if cls_res not in schema:print(f'The type model inferenced {cls_res} which is not in schema dict, exited.')exit()properties_str = ', '.join(schema[cls_res])schema_str_list = f'"{cls_res}"({properties_str})'sentence_with_ie_prompt = IE_PATTERN.format(sentence, schema_str_list)full_input_text = build_prompt(sentence_with_ie_prompt, custom_settings["ie_pre_history"])inputs = tokenizer(full_input_text, return_tensors="pt")inputs = inputs.to('cpu')output_sequences = model.generate(inputs["input_ids"], max_length=2000, attention_mask=inputs.attention_mask)generated_reply = tokenizer.decode(output_sequences[0], skip_special_tokens=True)print(generated_reply[len(full_input_text):])
3.6 测试模型
最后,我们使用一些测试数据来验证模型的效果。
sentences = ['“2024 年潮流双肩包,品牌 PQR,材质耐磨,内部空间大。定价399元,优惠后价格349元。周销量800个。”','2024 年智能手表,品牌为华为,功能强大,续航持久。售价1299元,便宜价999元,目前有黑色和银色可选。月销量3000只。'
]
custom_settings = init_prompts()
inference(sentences, custom_settings)
4. 结果展示
运行上述代码后,模型将输出从测试数据中抽取的信息。例如,对于第一句测试数据,模型可能会输出类似以下的结果:
5. 总结
本文介绍了一个基于 Hugging Face Transformers 和 PyTorch 的信息抽取实现。通过使用预训练的语言模型和提示模板,我们能够从文本中抽取特定的结构化信息。这种方法的优点是不需要大量的标注数据,并且能够处理复杂的文本结构。然而,它也有一些局限性,例如对模型的依赖较大,且生成的结果可能需要进一步验证和修正。
未来,我们可以探索更多优化方法,例如微调模型以提高其在特定任务上的性能,或者结合其他技术(如规则匹配)来提高抽取的准确性。此外,我们还可以尝试将这种方法应用于其他类型的信息抽取任务,以验证其通用性。