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

相关信息表格