Cuckoo C4

🚀 クックー 🐦 [Github]

このリポジトリには、論文 Cuckoo: An IE Free Rider Hatched by Massive Nutrition in LLM's Nest のモデルが含まれています。

クックーは、大規模言語モデルの次のトークン予測パラダイムを模倣した小規模（300M）の情報抽出（IE）モデルです。語彙からの検索ではなく、クックーは与えられた入力コンテキスト内でトークンにタグを付けることで次のトークンを予測します。下図のようにです。

クックーは、以前のIE事前学習とは大幅に異なり、あらゆるテキストリソースを利用して自身を強化することができ、特にLLM用に整理されたデータを無料で利用することができます！

現在、以下のデータで事前学習されたクックーのチェックポイントをオープンソースで公開しています。

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データセット、複数の選択肢QAデータセット、MMLU、SQuAD、DROP、MNLI、SNLI。(Cuckoo-C4-Super-Rainbow 🦸🌈🐦🛠️)

✨ 主な機能

🚀 性能デモ

クックーを使って、あらゆるIEタスクに対する想像を超えた適応効率を体験してください！

📦 インストール

本プロジェクトでは、事前学習済みのモデルを提供しています。以下の手順でモデルを利用することができます。

まず、必要なライブラリをインストールします。

pip install transformers torch spacy
python -m spacy download en_core_web_sm

💻 使用例

基本的な使用法

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)

高度な使用法

次のトークン抽出関数の定義

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

抽出関数の呼び出し

ケース1: 基本的なエンティティと関係の理解

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

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

実行すると、以下のような結果が得られます。

What are the people 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']
Which city does George live in? []

ここで、[] はクックーが抽出する次のトークンがないと判断したことを示しています。

ケース2: より長いコンテキスト

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']

ケース3: 知識クイズ

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)

実行すると、以下のような結果が得られます。

grass ['green']
sea ['blue']
fire ['red']
night []

これは、クックーが妥当なスパンを抽出していないが、コンテキストを理解する知識を持っていることを示しています。

📚 ドキュメント

ファイル情報

リポジトリには以下のファイル情報が含まれています。

📄 ライセンス

このプロジェクトはMITライセンスの下で公開されています。