基于Zero-shot实现LLM文本分类

2024-06-09 13:11:10 浏览数 (1)

基于Zero-shot方式实现LLM文本分类

项目介绍

  • 随着ChatGPT等大模型的出现,自然语言处理任务的方式也在逐渐发生改变。随着科技的进步和发展,未来如果大模型强大到一定程度,我们或许不再为每一个具体任务去finetune一个模型,而是使用同一个大模型,对于不同的任务进行prompt设计,以解决不同的任务问题。
  • 基于清华开源大模型ChatGLM-6B, 提供多个自然语言处理任务如何通过prompt来实现zero-shot

环境准备

本项目运行前请安装相关依赖包:

  • protobuf>=3.19.5
  • transformers>=4.26.1
  • icetk
  • cpm_kernels
  • streamlit==1.17.0

1 LLM文本分类任务介绍

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

代码语言:javascript复制
1. 加拿大(英语/法语:Canada),首都渥太华,位于北美洲北部。东临大西洋,西濒太平洋,西北部邻美国阿拉斯加州,南接美国本土,北靠北冰洋。气候大部分为亚寒带针叶林气候和湿润大陆性气候,北部极地区域为极地长寒气候。
2. 《琅琊榜》是由山东影视传媒集团、山东影视制作有限公司、北京儒意欣欣影业投资有限公司、北京和颂天地影视文化有限公司、北京圣基影业有限公司、东阳正午阳光影视有限公司联合出品,由孔笙、李雪执导,胡歌、刘涛、王凯、黄维德、陈龙、吴磊、高鑫等主演的古装剧。
3. 《满江红》是由张艺谋执导,沈腾、易烊千玺、张译、雷佳音、岳云鹏、王佳怡领衔主演,潘斌龙、余皑磊主演,郭京飞、欧豪友情出演,魏翔、张弛、黄炎特别出演,许静雅、蒋鹏宇、林博洋、飞凡、任思诺、陈永胜出演的悬疑喜剧电影。
4. 布宜诺斯艾利斯(Buenos Aires,华人常简称为布宜诺斯)是阿根廷共和国(the Republic of Argentina,República Argentina)的首都和最大城市,位于拉普拉塔河南岸、南美洲东南部、河对岸为乌拉圭东岸共和国。
5. 张译(原名张毅),1978年2月17日出生于黑龙江省哈尔滨市,中国内地男演员。1997年至2006年服役于北京军区政治部战友话剧团。2006年,主演军事励志题材电视剧《士兵突击》。

我们的目的是期望模型能够帮助我们识别出这5段话中,每一句话描述的是一个什么类型的物体。

这 5 段描述分别节选自:

代码语言:javascript复制
['加拿大', '《琅琊榜》', '《满江红》', '布宜诺斯艾利斯', '张译']

因此,我们期望模型输出的结果为:

代码语言:javascript复制
['国家', '电视剧', '电影', '城市', '人物']

2 Prompt设计

对于大模型来讲,prompt 的设计非常重要,一个 明确 的 prompt 能够帮助我们更好从大模型中获得我们想要的结果。

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

  • 需要向模型解释什么叫作「文本分类任务」
  • 需要让模型按照我们指定的格式输出

为了让模型知道什么叫做「文本分类」,我们借用 Incontext Learning 的方式,先给模型展示几个正确的例子:

代码语言:javascript复制
>>> User: "岳云鹏,本名岳龙刚,1985年4月15日出生于河南省濮阳市南乐县,中国内地相声、影视男演员。" 是 ['人物', '书籍', '电视剧', '电影', '城市', '国家'] 里的什么类别?
>>> Bot: 人物
>>> User: "《三体》是刘慈欣创作的长篇科幻小说系列,由《三体》《三体2:黑暗森林》《三体3:死神永生》组成,第一部于2006年5月起在《科幻世界》杂志上连载,第二部于2008年5月首次出版,第三部则于2010年11月出版。" 是 ['人物', '书籍', '电视剧', '电影', '城市', '国家'] 里的什么类别?
>>> Bot: 书籍

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

3 分类任务代码实现

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

本次分类任务实现的主要过程:

  • 构造prompt
  • 实现模型预测

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

llm_classification.py脚本中包含两个函数:init_prompts()和inference()

3.1 导入必备的工具包
代码语言:javascript复制
# —*-coding:utf-8-*-
"""
利用 LLM 进行文本分类任务。
"""
from rich import print
from rich.console import Console
from transformers import AutoTokenizer, AutoModel


# 提供所有类别以及每个类别下的样例
class_examples = {
        '人物': '岳云鹏,本名岳龙刚,1985年4月15日出生于河南省濮阳市南乐县,中国内地相声、影视男演员 [1]  。2005年,首次登台演出。2012年,主演卢卫国执导的喜剧电影《就是闹着玩的》。2013年在北京举办相声专场。',
        '书籍': '《三体》是刘慈欣创作的长篇科幻小说系列,由《三体》《三体2:黑暗森林》《三体3:死神永生》组成,第一部于2006年5月起在《科幻世界》杂志上连载,第二部于2008年5月首次出版,第三部则于2010年11月出版。',
        '电视剧': '《狂飙》是由中央电视台、爱奇艺出品,留白影视、中国长安出版传媒联合出品,中央政法委宣传教育局、中央政法委政法综治信息中心指导拍摄,徐纪周执导,张译、张颂文、李一桐、张志坚、吴刚领衔主演,倪大红、韩童生、李建义、石兆琪特邀主演,李健、高叶、王骁等主演的反黑刑侦剧。',
        '电影': '《流浪地球》是由郭帆执导,吴京特别出演、屈楚萧、赵今麦、李光洁、吴孟达等领衔主演的科幻冒险电影。影片根据刘慈欣的同名小说改编,故事背景设定在2075年,讲述了太阳即将毁灭,毁灭之后的太阳系已经不适合人类生存,而面对绝境,人类将开启“流浪地球”计划,试图带着地球一起逃离太阳系,寻找人类新家园的故事。',
        '城市': '乐山,古称嘉州,四川省辖地级市,位于四川中部,四川盆地西南部,地势西南高,东北低,属中亚热带气候带;辖4区、6县,代管1个县级市,全市总面积12720.03平方公里;截至2021年底,全市常住人口315.1万人。',
        '国家': '瑞士联邦(Swiss Confederation),简称“瑞士”,首都伯尔尼,位于欧洲中部,北与德国接壤,东临奥地利和列支敦士登,南临意大利,西临法国。地处北温带,四季分明,全国地势高峻,矿产资源匮乏,森林及水力资源丰富,总面积41284平方千米,全国由26个州组成(其中6个州为半州)。'
    }

3.2 构建init_prompts()函数
  • 目的:进行prompt设计
  • 具体代码实现:
代码语言:javascript复制
def init_prompts():
    """
    初始化前置prompt,便于模型做 incontext learning。
    """
    class_list = list(class_examples.keys())
    pre_history = [
        (
            f'现在你是一个文本分类器,你需要按照要求将我给你的句子分类到:{class_list}类别中。',
            f'好的。'
        )
    ]

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

3.3 构建inference()函数

  • 目的:实现文本分类
  • 具体代码实现
代码语言:javascript复制
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_prompt = f"“{sentence}”是 {custom_settings['class_list']} 里的什么类别?"
            response, history = model.chat(tokenizer, sentence_with_prompt, history=custom_settings['pre_history'])
        print(f'>>> [bold bright_red]sentence: {sentence}')
        print(f'>>> [bold bright_green]inference answer: {response}')
  • 代码调用
代码语言:javascript复制
if __name__ == '__main__':
    console = Console()

    device = 'cpu'
    tokenizer = AutoTokenizer.from_pretrained("./ChatGLM-6B/THUDM/chatglm-6b-int4", trust_remote_code=True)
    model = AutoModel.from_pretrained("./ChatGLM-6B/THUDM/chatglm-6b-int4", trust_remote_code=True).float()
    model.to(device)

    sentences = [
        "加拿大(英语/法语:Canada),首都渥太华,位于北美洲北部。东临大西洋,西濒太平洋,西北部邻美国阿拉斯加州,南接美国本土,北靠北冰洋。气候大部分为亚寒带针叶林气候和湿润大陆性气候,北部极地区域为极地长寒气候。",
        "《琅琊榜》是由山东影视传媒集团、山东影视制作有限公司、北京儒意欣欣影业投资有限公司、北京和颂天地影视文化有限公司、北京圣基影业有限公司、东阳正午阳光影视有限公司联合出品,由孔笙、李雪执导,胡歌、刘涛、王凯、黄维德、陈龙、吴磊、高鑫等主演的古装剧。",
        "《满江红》是由张艺谋执导,沈腾、易烊千玺、张译、雷佳音、岳云鹏、王佳怡领衔主演,潘斌龙、余皑磊主演,郭京飞、欧豪友情出演,魏翔、张弛、黄炎特别出演,许静雅、蒋鹏宇、林博洋、飞凡、任思诺、陈永胜出演的悬疑喜剧电影。",
        "布宜诺斯艾利斯(Buenos Aires,华人常简称为布宜诺斯)是阿根廷共和国(the Republic of Argentina,República Argentina)的首都和最大城市,位于拉普拉塔河南岸、南美洲东南部、河对岸为乌拉圭东岸共和国。",
        "张译(原名张毅),1978年2月17日出生于黑龙江省哈尔滨市,中国内地男演员。1997年至2006年服役于北京军区政治部战友话剧团。2006年,主演军事励志题材电视剧《士兵突击》。"
    ]
    
    custom_settings = init_prompts()
    print(custom_settings)

    inference(
        sentences,
        custom_settings
    )

小结

主要介绍了基于Zero-shot方式实现ChatGLM模型来完成文本分类的过程。

0 人点赞