Contra Bottleneck T5 Large Wikipedia

🚀 瓶颈T5模型 ⏳

瓶颈T5模型为我许多探索潜在空间中检查和编辑文本界面的实验与演示提供了支持。该模型是一个文本自动编码器，能够将最多512个标记的文本编码为嵌入向量，然后从嵌入向量中重构原始文本。该模型生成的嵌入空间结构还允许通过潜在空间中的向量运算对文本进行语义编辑。

🚀 快速开始

模型简介

使用该模型生成的嵌入向量，我们可以在文本片段之间进行语义插值，并利用句子的潜在属性（如长度、语气、结构或主题）对句子进行编辑。

所有瓶颈T5模型都在经过筛选的英文维基百科子集上进行训练，在编码和解码百科全书及其他类似类型的文本时表现最佳。技术含量高、对话式或其他非常规的文本可能超出了模型的分布范围，模型在处理此类输入时可能表现不佳。

瓶颈T5嵌入向量始终被归一化为长度为1；编码器生成长度为1的嵌入向量，解码器的任何输入也将被归一化为长度为1。

模型信息

💻 使用示例

基础用法

该模型目前处于原型阶段，基于T5语言模型实现，因此我们需要一个小包装类来使用它进行文本嵌入和生成：

import os
import torch
import torch.nn as nn
import torch.nn.functional as F

from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM

class BottleneckT5Autoencoder:
    def __init__(self, model_path: str, device='cpu'):
        self.device = device
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, model_max_length=512)
        self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).to(self.device)
        self.model.eval()

    @torch.no_grad()
    def embed(self, text: str) -> torch.FloatTensor:
        inputs = self.tokenizer(text, return_tensors='pt').to(self.device)
        decoder_inputs = self.tokenizer('', return_tensors='pt').to(self.device)
        return self.model(
            **inputs,
            decoder_input_ids=decoder_inputs['input_ids'],
            encode_only=True,
        )[0]

    @torch.no_grad()
    def generate_from_latent(self, latent: torch.FloatTensor, max_length=512, temperature=1.0) -> str:
        dummy_text = '.'
        dummy = self.embed(dummy_text)
        perturb_vector = latent - dummy
        self.model.perturb_vector = perturb_vector
        input_ids = self.tokenizer(dummy_text, return_tensors='pt').to(self.device).input_ids
        output = self.model.generate(
            input_ids=input_ids,
            max_length=max_length,
            do_sample=True,
            temperature=temperature,
            top_p=0.9,
            num_return_sequences=1,
        )
        return self.tokenizer.decode(output[0], skip_special_tokens=True)

然后我们可以基于模型类初始化这个自动编码器类：

device = 'cuda' if torch.cuda.is_available() else 'cpu'
autoencoder = BottleneckT5Autoencoder(model_path='thesephist/contra-bottleneck-t5-large-wikipedia', device=device)

使用 .embed(text: str) 和 .generate_from_latent(embedding: torch.FloatTensor) 对文本进行嵌入和反嵌入：

texts = [
    'The quick brown fox jumps over the lazy dog',
    'Hi there! My name is Linus, and I spend a lot of my time thinking about latent spaces of neural network models.',
    'Notion is a single space where you can think, write, and plan. Capture thoughts, manage projects, or even run an entire company — and do it exactly the way you want.',
]

for t in texts:
    embedding = autoencoder.embed(t)
    reconstruction = autoencoder.generate_from_latent(embedding)
    print(reconstruction)

上述代码将输出以下文本：

The quick brown fox jumps over the lazy dog
I'm named after Linus, and I spend a lot of my time thinking about neural networks of latent space models.
Notion is a single place where you can think, plan, and spend time. Capture ideas, manage projects, and even do your own writing — or plan it exactly the way you want.

更多示例

有关如何使用该模型通过Contra进行插值和语义编辑的更多示例，请参阅 此Google Colab笔记本。

🔧 技术细节

训练详情

Contra模型从 适应语言模型的T5 v1.1检查点 初始化，并在经过长度筛选的英文 维基百科 数据集子集上进行了一个轮次的训练，作为去噪自动编码器，随机屏蔽30%的标记，使用Adafactor优化器。

模型家族和检查点

我建议首先使用 thesephist/contra-bottleneck-t5-large-wikipedia 进行实验，它在模型大小和输出质量之间取得了很好的平衡。不过，我已经训练了四个参数范围从3.3亿到30亿的变体：

• thesephist/contra-bottleneck-t5-small-wikipedia：6000万参数，512个嵌入维度
• thesephist/contra-bottleneck-t5-base-wikipedia：2.2亿参数，768个嵌入维度
• thesephist/contra-bottleneck-t5-large-wikipedia：7.7亿参数，1024个嵌入维度
• thesephist/contra-bottleneck-t5-xl-wikipedia：30亿参数，2048个嵌入维度

📄 许可证

本模型使用MIT许可证。