Rankgen T5 Xl All

🚀 RankGen

RankGen是一套編碼器模型（參數範圍為1億 - 12億），它可以將任何預訓練英文語言模型的前綴和生成內容映射到一個共享向量空間。RankGen可用於對語言模型生成的多個完整樣本進行重新排序，還能作為評分函數融入束搜索，顯著提升生成質量（MAUVE指標從0.77提升到0.85，根據英文作家人工標註結果，偏好度達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 來加載RankGen，它是一個圍繞HuggingFace API的小包裝器，可自動正確預處理數據和進行分詞。你可以下載我們的倉庫並安裝API，或者從下面複製實現代碼。

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
        }

📄 許可證

本項目採用Apache-2.0許可證。

主要倉庫

https://github.com/martiansideofthemoon/rankgen

相關信息表格