OpenPrompt:一个开源的提示学习框架

OpenPrompt:开启提示学习新范式

在自然语言处理(NLP)领域,提示学习(Prompt Learning)作为一种新兴的范式正在迅速崛起。它通过设计文本模板来改变输入,直接利用预训练语言模型(PLMs)进行下游任务,展现出了强大的潜力。为了推动这一技术的发展和应用,清华大学自然语言处理实验室(THUNLP)开发了OpenPrompt - 一个开源的提示学习框架。

什么是OpenPrompt?

OpenPrompt是一个基于PyTorch构建的标准化、灵活且可扩展的框架,旨在为提示学习提供全面的实现方案。它具有以下主要特点:

支持多种预训练语言模型:OpenPrompt可以直接加载Hugging Face Transformers库中的PLMs,未来还将支持其他库实现的模型。
提供丰富的提示学习方法:包括模板构建、标签词映射、优化策略等多个方面的实现。
模块化设计:框架的核心组件如Template、Verbalizer等都采用模块化设计,便于用户自定义和扩展。
易用性强:提供了直观的API和详细的教程,使用户能够快速上手。
研究友好:支持灵活组合不同的PLMs、任务格式和提示模块,有利于进行对比实验。

OpenPrompt框架概览

OpenPrompt的核心组件

OpenPrompt框架主要包含以下几个核心组件:

PromptModel:整合PLM、Template和Verbalizer的核心类,用于训练和推理。
Template:负责将原始输入文本包装成提示模板的类。OpenPrompt提供了多种模板实现,如ManualTemplate、PrefixTemplate等。
Verbalizer:构建标签和目标词之间映射关系的类。包括ManualVerbalizer、AutomaticVerbalizer等多种实现。
PromptDataLoader:提示学习版本的数据加载器,集成了Tokenizer、Template等组件。
Trainer:封装了训练和评估流程的类,提供了分类、生成等多种任务的实现。

通过组合这些灵活的组件,用户可以快速构建出适合自己任务的提示学习pipeline。

使用OpenPrompt的简单示例

下面通过一个简单的情感分析任务,展示如何使用OpenPrompt:

定义任务:

from openprompt.data_utils import InputExample

classes = ["negative", "positive"]
dataset = [
    InputExample(guid=0, text_a="Albert Einstein was one of the greatest intellects of his time."),
    InputExample(guid=1, text_a="The film was badly made."),
]

加载预训练模型:

from openprompt.plms import load_plm
plm, tokenizer, model_config, WrapperClass = load_plm("bert", "bert-base-cased")

定义模板:

from openprompt.prompts import ManualTemplate
promptTemplate = ManualTemplate(
    text = '{"placeholder":"text_a"} It was {"mask"}',
    tokenizer = tokenizer,
)

定义Verbalizer:

from openprompt.prompts import ManualVerbalizer
promptVerbalizer = ManualVerbalizer(
    classes = classes,
    label_words = {
        "negative": ["bad"],
        "positive": ["good", "wonderful", "great"],
    },
    tokenizer = tokenizer,
)

构建PromptModel:

from openprompt import PromptForClassification
promptModel = PromptForClassification(
    template = promptTemplate,
    plm = plm,
    verbalizer = promptVerbalizer,
)

创建数据加载器:

from openprompt import PromptDataLoader
data_loader = PromptDataLoader(
    dataset = dataset,
    tokenizer = tokenizer, 
    template = promptTemplate,
    tokenizer_wrapper_class=WrapperClass,
)

进行推理:

import torch

promptModel.eval()
with torch.no_grad():
    for batch in data_loader:
        logits = promptModel(batch)
        preds = torch.argmax(logits, dim = -1)
        print(classes[preds])