Cuckoo C4 Super Rainbow

🚀 布谷鸟模型（Cuckoo）

布谷鸟（Cuckoo）是一个小型（3亿参数）的信息提取（IE）模型，它模仿大语言模型的下一个标记预测范式。该模型通过在给定的输入上下文中标记来预测下一个标记，而不是从词汇表中检索。这使得它能够利用任何文本资源来提升自身性能，尤其可以借助为大语言模型整理的数据实现高效学习。

🚀 快速开始

快速体验下一个标记提取

我们推荐使用最强的超级彩虹布谷鸟模型（Cuckoo-C4-Super-Rainbow）进行零样本提取。

1️⃣ 首先加载模型和分词器

from transformers import AutoModelForTokenClassification, AutoTokenizer
import torch
import spacy

nlp = spacy.load("en_core_web_sm")

device = torch.device("cuda:0")
path = f"KomeijiForce/Cuckoo-C4-Super-Rainbow"
tokenizer = AutoTokenizer.from_pretrained(path)
tagger = AutoModelForTokenClassification.from_pretrained(path).to(device)

2️⃣ 定义下一个标记提取函数

def next_tokens_extraction(text):

    def find_sequences(lst):
        sequences = []
        i = 0
        while i < len(lst):
            if lst[i] == 0:
                start = i
                end = i
                i += 1
                while i < len(lst) and lst[i] == 1:
                    end = i
                    i += 1
                sequences.append((start, end+1))
            else:
                i += 1
        return sequences

    text = " ".join([token.text for token in nlp(text)])

    inputs = tokenizer(text, return_tensors="pt").to(device)
    tag_predictions = tagger(**inputs).logits[0].argmax(-1)

    predictions = [tokenizer.decode(inputs.input_ids[0, seq[0]:seq[1]]).strip() for seq in find_sequences(tag_predictions)]
    
    return predictions

3️⃣ 调用函数进行提取！

基础用法

text = "Tom and Jack went to their trip in Paris."

for question in [
    "What is the person mentioned here?",
    "What is the city mentioned here?",
    "Who goes with Tom together?",
    "What do Tom and Jack go to Paris for?",
    "Where does George live in?",
]:
    prompt = f"User:\n\n{text}\n\nQuestion: {question}\n\nAssistant:"
    predictions = next_tokens_extraction(prompt)
    print(question, predictions)

你将得到类似如下的结果：

What is the person mentioned here? ['Tom', 'Jack']
What is the city mentioned here? ['Paris']
Who goes with Tom together? ['Jack']
What do Tom and Jack go to Paris for? ['trip']
Where does George live in? []

其中 [] 表示布谷鸟模型认为没有可提取的下一个标记。

高级用法

passage = f'''Ludwig van Beethoven (17 December 1770 – 26 March 1827) was a German composer and pianist. He is one of the most revered figures in the history of Western music; his works rank among the most performed of the classical music repertoire and span the transition from the Classical period to the Romantic era in classical music. His early period, during which he forged his craft, is typically considered to have lasted until 1802. From 1802 to around 1812, his middle period showed an individual development from the styles of Joseph Haydn and Wolfgang Amadeus Mozart, and is sometimes characterised as heroic. During this time, Beethoven began to grow increasingly deaf. In his late period, from 1812 to 1827, he extended his innovations in musical form and expression.'''

for question in [
    "What are the people mentioned here?",
    "What is the job of Beethoven?",
    "How famous is Beethoven?",
    "When did Beethoven's middle period showed an individual development?",
]:
    text = f"User:\n\n{passage}\n\nQuestion: {question}\n\nAssistant:"
    predictions = next_tokens_extraction(text)
    print(question, predictions)

你将得到类似如下的结果：

What are the people mentioned here? ['Ludwig van Beethoven', 'Joseph Haydn', 'Wolfgang Amadeus Mozart']
What is the job of Beethoven? ['composer and pianist']
How famous is Beethoven? ['one of the most revered figures in the history of Western music']
When did Beethoven's middle period showed an individual development? ['1802']

小样本适配

布谷鸟模型在小样本适配自身任务方面表现出色。以下是一些示例：

• 以CoNLL2003为例：运行命令 bash run_downstream.sh conll2003.5shot KomeijiForce/Cuckoo-C4-Rainbow，你将在 models/cuckoo-conll2003.5shot 中得到一个微调后的模型。然后可以使用脚本 python eval_conll2003.py 对模型进行基准测试，其F1性能约为80。
• 机器阅读理解（SQuAD）适配：运行命令 bash run_downstream.sh squad.32shot KomeijiForce/Cuckoo-C4-Rainbow，你将在 models/cuckoo-squad.32shot 中得到一个微调后的模型。然后使用脚本 python eval_squad.py 进行基准测试，F1性能约为88。

如果你要微调自己的任务，需要创建一个Jsonlines文件，每行包含 {"words": [...], "ner": [...]}，例如：

{"words": ["I", "am", "John", "Smith", ".", "Person", ":"], "ner": ["O", "O", "B", "I", "O", "O", "O"]}

这表示 "John Smith" 被预测为下一个标记。

你可以参考以下提示开始构建自己的下游数据集：

构建好自己的下游数据集后，将其保存为 my_downstream.json，然后运行命令 bash run_downstream.sh my_downstream KomeijiForce/Cuckoo-C4-Rainbow。你将在 models/cuckoo-my_downstream 中找到适配后的布谷鸟模型。

训练自己的布谷鸟模型

我们在 nte_data_collection.py 文件中提供了将文本转换为NTE实例的脚本，以C4为例，转换结果可以在 cuckoo.c4.example.json 中查看。该脚本易于适配其他资源，如实体、查询和问题，你可以将自己的数据修改为NTE格式来训练自己的布谷鸟模型！运行 run_cuckoo.sh 脚本进行示例预训练：

python run_ner.py \
  --model_name_or_path roberta-large \
  --train_file cuckoo.c4.example.json \
  --output_dir models/cuckoo-c4-example \
  --per_device_train_batch_size 4\
  --gradient_accumulation_steps 16\
  --num_train_epochs 1\
  --save_steps 1000\
  --learning_rate 0.00001\
  --do_train \
  --overwrite_output_dir

你将在 models/cuckoo-c4-example 中得到一个示例布谷鸟模型。如果预训练数据过少，模型性能可能不佳。你可以调整 nte_data_collection.py 中的超参数或修改转换逻辑，以实现更好的预训练性能。

✨ 主要特性

• 创新的信息提取范式：布谷鸟模型模仿大语言模型的下一个标记预测范式，通过在给定输入上下文中标记来预测下一个标记，与传统的信息提取预训练方法有很大不同。
• 数据利用高效：可以利用任何文本资源来提升自身性能，尤其能够借助为大语言模型整理的数据实现高效学习。
• 多场景适配能力：支持零样本提取和小样本适配，能够快速适应不同的信息提取任务。

📦 安装指南

文档未提供具体安装步骤，暂不展示。

💻 使用示例

上述快速开始部分已包含详细的使用示例，此处不再赘述。

📚 详细文档

性能展示

布谷鸟模型在多种信息提取任务上表现出色，以下是其与其他模型的性能对比：

可用的预训练模型

目前开源的布谷鸟模型检查点基于以下数据进行预训练：

1. 从C4转换而来的1亿个下一个标记提取（NTE）实例。（Cuckoo-C4 🐦）
2. Cuckoo-C4 + 从有监督微调数据集TuluV3转换而来的260万个下一个标记提取（NTE）实例。（Cuckoo-C4-Instruct 🐦🛠️）
3. Cuckoo-C4-Instruct + MultiNERD、MetaIE、NuNER、MRQA（不包括SQuAD、DROP）。（Cuckoo-C4-Rainbow 🌈🐦🛠️）
4. Cuckoo-C4-Rainbow + 多个NER数据集、WizardLM数据集、多项选择问答数据集、MMLU、SQuAD、DROP、MNLI、SNLI。（Cuckoo-C4-Super-Rainbow 🦸🌈🐦🛠️）

🔧 技术细节

文档未提供具体技术细节，暂不展示。

📄 许可证

本项目使用Apache-2.0许可证。

🐾 引用

@article{DBLP:journals/corr/abs-2502-11275,
  author       = {Letian Peng and
                  Zilong Wang and
                  Feng Yao and
                  Jingbo Shang},
  title        = {Cuckoo: An {IE} Free Rider Hatched by Massive Nutrition in {LLM}'s Nest},
  journal      = {CoRR},
  volume       = {abs/2502.11275},
  year         = {2025},
  url          = {https://doi.org/10.48550/arXiv.2502.11275},
  doi          = {10.48550/arXiv.2502.11275},
  eprinttype   = {arXiv},
  eprint       = {2502.11275},
  timestamp    = {Mon, 17 Feb 2025 19:32:20 +0000},
  biburl       = {https://dblp.org/rec/journals/corr/abs-2502-11275.bib},
  bibsource    = {dblp computer science bibliography, https://dblp.org}
}