基于Zero-shot实现LLM信息抽取

@小森

发布于 2024-06-09 13:12:27

850

发布于 2024-06-09 13:12:27

文章被收录于专栏：xiaosenxiaosen

基于Zero-shot方式实现LLM信息抽取

在当今这个信息爆炸的时代，从海量的文本数据中高效地抽取关键信息显得尤为重要。随着自然语言处理（NLP）技术的不断进步，信息抽取任务也迎来了新的突破。近年来，基于Zero-shot（零样本学习）的大型语言模型（LLM）在信息抽取领域展现出了强大的潜力。这种方法能够在没有预先标注数据的情况下，通过理解自然语言指令来完成信息抽取任务，极大地提高了信息处理的灵活性和效率。

1 LLM信息抽取任务介绍

首先，我们定义信息抽取的Schema：

# 定义不同实体下的具备属性
schema = {
    '人物': ['姓名', '性别', '出生日期', '出生地点', '职业', '获得奖项', '实体类型'],
    '书籍': ['作者', '类型', '发行时间', '定价', '实体类型'],
    '电视剧': ['导演', '演员', '题材', '出品方', '实体类型']
}

下面几段文本来自百度百科：

1. 张译（原名张毅），1978年2月17日出生于黑龙江省哈尔滨市，中国内地男演员。1997年至2006年服役于北京军区政治部战友话剧团。2006年，主演军事励志题材电视剧《士兵突击》。
2. 《琅琊榜》是由山东影视传媒集团、山东影视制作有限公司、北京儒意欣欣影业投资有限公司、北京和颂天地影视文化有限公司、北京圣基影业有限公司、东阳正午阳光影视有限公司联合出品，由孔笙、李雪执导，胡歌、刘涛、王凯、黄维德、陈龙、吴磊、高鑫等主演的古装剧。

我们的目的是期望模型能够帮助我们识别出这2段话中的SPO三元组信息。

2 Prompt设计

在该任务的 prompt 设计中，我们主要考虑 2 点：

需要向模型解释什么叫作「信息抽取任务」
需要让模型按照我们指定的格式（json）输出

为了让模型知道什么叫做「信息抽取」，我们借用 Incontext Learning 的方式，先给模型展示几个正确的例子：

>>> User: 岳云鹏，本名岳龙刚，1985年4月15日出生于河南省濮阳市南乐县，中国内地相声、影视男演员。2005年，首次登台演出。2012年，主演卢卫国执导的喜剧电影《就是闹着玩的》。2013年在北京举办相声专场。提取上述句子中“人物”(姓名, 性别, 出生日期, 出生地点, 职业, 获得奖项)类型的实体，并按照JSON格式输出，上述句子中没有的信息用['原文中未提及']来表示，多个值之间用','分隔。
>>> Bot: {"姓名": ["岳云鹏"], "性别": ["男"], "出生日期": ["1985年4月15日"], "出生地点": ["河南省濮阳市南乐县"], "职业": ["相声演员", "影视演员"], "获得奖项": ["原文中未提及"]}
...

其中，User 代表我们输入给模型的句子，Bot 代表模型的回复内容。注意：上述例子中 Bot 的部分也是由人工输入的，其目的是希望看到在看到类似 User 中的句子时，模型应当做出类似 Bot 的回答。

3 关系抽取任务代码实现

本章节使用的模型为ChatGLM-6B，参数参数较大（6B），下载到本地大概需要 12G+ 的磁盘空间，请确保磁盘有充足的空间。此外，加载模型大概需要 13G 左右的显存，如果您显存不够，可以进行模型量化加载以缩小模型成本。

本次信息抽取任务实现的主要过程：

构造prompt
先对句子做分类
再进行信息抽取

代码存放位置：/Users/**/PycharmProjects/llm/zero-shot/llm_information_extraction.py

llm_information_extraction.py脚本中包含三个函数：init_prompts()、clean_response()和inference()

3.1 导入必备的工具包

"""
利用 LLM 进行信息抽取任务，先对句子做分类，再进行信息提取。
"""
import re
import json
import os


from rich import print
from rich.console import Console
from transformers import AutoTokenizer, AutoModel



# 分类 example
class_examples = {
        '人物': '岳云鹏，本名岳龙刚，1985年4月15日出生于河南省濮阳市南乐县，中国内地相声、影视男演员。2005年，首次登台演出。2012年，主演卢卫国执导的喜剧电影《就是闹着玩的》。2013年在北京举办相声专场。',
        '书籍': '《三体》是刘慈欣创作的长篇科幻小说系列，由《三体》《三体2：黑暗森林》《三体3：死神永生》组成，第一部于2006年5月起在《科幻世界》杂志上连载，第二部于2008年5月首次出版，第三部则于2010年11月出版。',
        '电视剧': '《狂飙》是由中央电视台、爱奇艺出品，留白影视、中国长安出版传媒联合出品，中央政法委宣传教育局、中央政法委政法综治信息中心指导拍摄，徐纪周执导，张译、张颂文、李一桐、张志坚、吴刚领衔主演，倪大红、韩童生、李建义、石兆琪特邀主演，李健、高叶、王骁等主演的反黑刑侦剧。',
    }
class_list = list(class_examples.keys())

CLS_PATTERN = f"“{{}}”是 {class_list} 里的什么类别？"


# 定义不同实体下的具备属性
schema = {
    '人物': ['姓名', '性别', '出生日期', '出生地点', '职业', '获得奖项'],
    '书籍': ['书名', '作者', '类型', '发行时间', '定价'],
    '电视剧': ['电视剧名称', '导演', '演员', '题材', '出品方']
}

IE_PATTERN = "{}\n\n提取上述句子中{}类型的实体，并按照JSON格式输出，上述句子中不存在的信息用['原文中未提及']来表示，多个值之间用','分隔。"


# 提供一些例子供模型参考
ie_examples = {
        '人物': [
                    {
                        'content': '岳云鹏，本名岳龙刚，1985年4月15日出生于河南省濮阳市南乐县，中国内地相声、影视男演员。',
                        'answers': {
                                        '姓名': ['岳云鹏'],
                                        '性别': ['男'],
                                        '出生日期': ['1985年4月15日'],
                                        '出生地点': ['河南省濮阳市南乐县'],
                                        '职业': ['相声演员', '影视演员'],
                                        '获得奖项': ['原文中未提及']
                            }
                    }
        ],
        '书籍': [
                    {
                        'content': '《三体》是刘慈欣创作的长篇科幻小说系列，由《三体》《三体2：黑暗森林》《三体3：死神永生》组成，第一部于2006年5月起在《科幻世界》杂志上连载，第二部于2008年5月首次出版，第三部则于2010年11月出版。',
                        'answers': {
                                        '书名': ['《三体》'],
                                        '作者': ['刘慈欣'],
                                        '类型': ['长篇科幻小说'],
                                        '发行时间': ['2006年5月','2008年5月','2010年11月'],
                                        '定价': ['原文中未提及']
                            }
                    }
        ]
}

3.2 构建init_prompts()函数

目的：进行prompt设计
具体代码实现：

def init_prompts():
    """
    初始化前置prompt，便于模型做 incontext learning。
    """
    class_list = list(class_examples.keys())
    cls_pre_history = [
        (
            f'现在你是一个文本分类器，你需要按照要求将我给你的句子分类到：{class_list}类别中。',
            f'好的。'
        )
    ]

    for _type, exmpale in class_examples.items():
        cls_pre_history.append((f'“{exmpale}”是 {class_list} 里的什么类别？', _type))

    ie_pre_history = [
        (
            "现在你需要帮助我完成信息抽取任务，当我给你一个句子时，你需要帮我抽取出句子中三元组，并按照JSON的格式输出，上述句子中没有的信息用['原文中未提及']来表示，多个值之间用','分隔。",
            '好的，请输入您的句子。'
        )
    ]

    for _type, example_list in ie_examples.items():
        for example in example_list:
            sentence = example['content']
            properties_str = ', '.join(schema[_type])
            schema_str_list = f'“{_type}”({properties_str})'
            sentence_with_prompt = IE_PATTERN.format(sentence, schema_str_list)
            ie_pre_history.append((
                f'{sentence_with_prompt}',
                f"{json.dumps(example['answers'], ensure_ascii=False)}"
            ))

    return {'ie_pre_history': ie_pre_history, 'cls_pre_history': cls_pre_history}

3.3 构建clean_response()函数

目的：模型结果后处理
具体代码实现

def clean_response(response: str):
    """
    后处理模型输出。

    Args:
        response (str): _description_
    """
    if '```json' in response:
        res = re.findall(r'```json(.*?)```', response)
        if len(res) and res[0]:
            response = res[0]
        response.replace('、', ',')
    try:
        return json.loads(response)
    except:
        return response

3.4 构建inference()函数

目的：模型实现信息抽取
具体代码实现

def inference(
        sentences: list,
        custom_settings: dict
    ):
    """
    推理函数。

    Args:
        sentences (List[str]): 待抽取的句子。
        custom_settings (dict): 初始设定，包含人为给定的 few-shot example。
    """
    for sentence in sentences:
        with console.status("[bold bright_green] Model Inference..."):
            sentence_with_cls_prompt = CLS_PATTERN.format(sentence)
            print(sentence_with_cls_prompt)
            cls_res, _ = model.chat(tokenizer, sentence_with_cls_prompt,  history=custom_settings['cls_pre_history'])


            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)
            ie_res, _ = model.chat(tokenizer, sentence_with_ie_prompt, history=custom_settings['ie_pre_history'])
            ie_res = clean_response(ie_res)
        print(f'>>> [bold bright_red]sentence: {sentence}')
        print(f'>>> [bold bright_green]inference answer: ')
        print(ie_res)

代码调用

if __name__ == '__main__':
    console = Console()
    #
    device = 'mps'
    tokenizer = AutoTokenizer.from_pretrained("/Users/ligang/PycharmProjects/llm/ChatGLM-6B/THUDM/chatglm-6b", trust_remote_code=True)
    model = AutoModel.from_pretrained("/Users/ligang/PycharmProjects/llm/ChatGLM-6B/THUDM/chatglm-6b", trust_remote_code=True).float()
    model.to(device)

    sentences = [
        '张译（原名张毅），1978年2月17日出生于黑龙江省哈尔滨市，中国内地男演员。1997年至2006年服役于北京军区政治部战友话剧团。2006年，主演军事励志题材电视剧《士兵突击》。',
        '《琅琊榜》是由山东影视传媒集团、山东影视制作有限公司、北京儒意欣欣影业投资有限公司、北京和颂天地影视文化有限公司、北京圣基影业有限公司、东阳正午阳光影视有限公司联合出品，由孔笙、李雪执导，胡歌、刘涛、王凯、黄维德、陈龙、吴磊、高鑫等主演的古装剧。',
    ]

    custom_settings = init_prompts()
    # print(f'custom_settings-->{custom_settings}')
    inference(
        sentences,
        custom_settings
    )