Cuckoo C4 Rainbow

🚀 布谷鸟模型（Cuckoo）🐦

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

🚀 快速开始

目前，我们开源了在不同数据集上预训练的布谷鸟模型检查点：

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 🦸🌈🐦🛠️）

✨ 主要特性

• 模仿预测范式：模仿大语言模型的下一个标记预测范式，通过在输入上下文中标记下一个标记来进行预测。
• 数据利用高效：可以使用任何文本资源来增强自身，特别是借助为大语言模型整理的数据。
• 多场景适应性：在多种信息提取任务中表现出色，支持少样本适应。

📊 性能展示

以下是布谷鸟模型在不同数据集上的性能表现，与其他模型进行了对比：

💻 使用示例

基础用法

我们推荐使用最强的超级彩虹布谷鸟模型（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)

• 更长上下文处理

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)

• 知识问答

for obj in ["grass", "sea", "fire", "night"]:
    text = f"User:\n\nChoices:\nred\nblue\ngreen.\n\nQuestion: What is the color of the {obj}?\n\nAssistant:\n\nAnswer:"
    predictions = next_tokens_extraction(text)
    print(obj, predictions)

🎯 少样本适应

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

• 以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"]}

创建好自己的下游数据集后，将其保存为 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 中的超参数或修改转换逻辑以适应自己的资源，从而获得更好的预训练性能。

🐾 引用

@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}
}

📄 许可证

本项目采用Apache 2.0许可证。