Rankgen T5 Xl All

🚀 RankGen

RankGenは、任意の事前学習済み英語言語モデルからの接頭辞と生成結果を共有ベクトル空間にマッピングするエンコーダモデルのセット（パラメータ数100M - 1.2B）です。このモデルは、言語モデルからの複数の全文サンプルを再ランク付けするために使用でき、ビームサーチにスコアリング関数として組み込むことで、生成品質を大幅に向上させることができます（MAUVEで0.85対0.77、英語の作家である人間のアノテーターによると75%の選好率）。また、密なリトリーバーのように使用することもでき、文学検索で最先端の性能を達成します。

🚀 クイックスタート

RankGenを使用するには、まずセットアップを行い、その後モデルをロードして使用します。以下に詳細な手順を説明します。

✨ 主な機能

• 任意の事前学習済み英語言語モデルからの接頭辞と生成結果を共有ベクトル空間にマッピング
• 言語モデルからの複数の全文サンプルを再ランク付け
• ビームサーチにスコアリング関数として組み込み、生成品質を大幅に向上
• 密なリトリーバーとして使用し、文学検索で最先端の性能を達成

📦 インストール

必要条件

（pipがこれらの依存関係を自動的にインストールします） Python 3.7+、torch（CUDA推奨）、transformers

インストール手順

python3.7 -m virtualenv rankgen-venv
source rankgen-venv/bin/activate
pip install rankgen

データはこちらから取得し、ルートディレクトリに配置してください。または、以下のようにgdownを使用することもできます。

gdown --folder https://drive.google.com/drive/folders/1DRG2ess7fK3apfB-6KoHb_azMuHbsIv4

RankGenチェックポイントが正しく読み込まれていることを確認するために、テストスクリプトを実行します。

python -m rankgen.test_rankgen_encoder --model_path kalpeshk2011/rankgen-t5-base-all

### 期待される出力
0.0009239262409127233
0.0011521980725477804

💻 使用例

基本的な使用法

RankGenEncoderを使用したモデルのロード

from rankgen import RankGenEncoder, RankGenGenerator

rankgen_encoder = RankGenEncoder("kalpeshk2011/rankgen-t5-xl-all")

# ベクトルのエンコード
prefix_vectors = rankgen_encoder.encode(["This is a prefix sentence."], vectors_type="prefix")
suffix_vectors = rankgen_encoder.encode(["This is a suffix sentence."], vectors_type="suffix")

# テキストの生成
# HuggingFace互換の言語モデルを使用
generator = RankGenGenerator(rankgen_encoder=rankgen_encoder, language_model="gpt2-medium")

inputs = ["Whatever might be the nature of the tragedy it would be over with long before this, and those moving black spots away yonder to the west, that he had discerned from the bluff, were undoubtedly the departing raiders. There was nothing left for Keith to do except determine the fate of the unfortunates, and give their bodies decent burial. That any had escaped, or yet lived, was altogether unlikely, unless, perchance, women had been in the party, in which case they would have been borne away prisoners."]

# ベースラインの核サンプリング
print(generator.generate_single(inputs, top_p=0.9)[0][0])
# 過剰生成と再ランク付け
print(generator.overgenerate_rerank(inputs, top_p=0.9, num_samples=10)[0][0])
# ビームサーチ
print(generator.beam_search(inputs, top_p=0.9, num_samples=10, beam_size=2)[0][0])

HuggingFace APIを使用したモデルのロード

from transformers import T5Tokenizer, AutoModel

tokenizer = T5Tokenizer.from_pretrained(f"google/t5-v1_1-xl")
model = AutoModel.from_pretrained("kalpeshk2011/rankgen-t5-xl-all", trust_remote_code=True)

高度な使用法

RankGenEncoderの実装

import tqdm
from transformers import T5Tokenizer, T5EncoderModel, AutoModel

class RankGenEncoder():
    def __init__(self, model_path, max_batch_size=32, model_size=None, cache_dir=None):
        assert model_path in ["kalpeshk2011/rankgen-t5-xl-all", "kalpeshk2011/rankgen-t5-xl-pg19", "kalpeshk2011/rankgen-t5-base-all", "kalpeshk2011/rankgen-t5-large-all"]
        self.max_batch_size = max_batch_size
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        if model_size is None:
            if "t5-large" in model_path or "t5_large" in model_path:
                self.model_size = "large"
            elif "t5-xl" in model_path or "t5_xl" in model_path:
                self.model_size = "xl"
            else:
                self.model_size = "base"
        else:
            self.model_size = model_size

        self.tokenizer = T5Tokenizer.from_pretrained(f"google/t5-v1_1-{self.model_size}", cache_dir=cache_dir)
        self.model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
        self.model.to(self.device)
        self.model.eval()

    def encode(self, inputs, vectors_type="prefix", verbose=False, return_input_ids=False):
        tokenizer = self.tokenizer
        max_batch_size = self.max_batch_size
        if isinstance(inputs, str):
            inputs = [inputs]
        if vectors_type == 'prefix':
            inputs = ['pre ' + input for input in inputs]
            max_length = 512
        else:
            inputs = ['suffi ' + input for input in inputs]
            max_length = 128

        all_embeddings = []
        all_input_ids = []
        for i in tqdm.tqdm(range(0, len(inputs), max_batch_size), total=(len(inputs) // max_batch_size) + 1, disable=not verbose, desc=f"Encoding {vectors_type} inputs:"):
            tokenized_inputs = tokenizer(inputs[i:i + max_batch_size], return_tensors="pt", padding=True)
            for k, v in tokenized_inputs.items():
                tokenized_inputs[k] = v[:, :max_length]
            tokenized_inputs = tokenized_inputs.to(self.device)
            with torch.inference_mode():
                batch_embeddings = self.model(**tokenized_inputs)
            all_embeddings.append(batch_embeddings)
            if return_input_ids:
                all_input_ids.extend(tokenized_inputs.input_ids.cpu().tolist())
        return {
            "embeddings": torch.cat(all_embeddings, dim=0),
            "input_ids": all_input_ids
        }

📚 ドキュメント

メインリポジトリ

https://github.com/martiansideofthemoon/rankgen

📄 ライセンス

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

情報一覧