Contra Bottleneck T5 Base Wikipedia

🚀 瓶颈T5模型 ⏳

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

🚀 快速开始

模型初始化

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)

文本编码与解码示例

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.

更多关于如何使用该模型进行插值和语义编辑的示例，请参考 这个Google Colab笔记本。

✨ 主要特性

• 文本自动编码与解码：能够将最多512个标记的文本编码为嵌入向量，并从嵌入向量中重构原始文本。
• 潜在空间语义编辑：通过潜在空间中的向量运算对文本进行语义编辑，如语义插值、根据长度、语气、结构或主题等潜在属性编辑句子。

📦 安装指南

文档未提及安装步骤，可参考相关依赖库（如transformers、torch等）的官方安装指南进行安装。

💻 使用示例

基础用法

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)

text = 'The quick brown fox jumps over the lazy dog'
embedding = autoencoder.embed(text)
reconstruction = autoencoder.generate_from_latent(embedding)
print(reconstruction)

高级用法

可参考 这个Google Colab笔记本 中的示例，使用该模型进行插值和语义编辑。

📚 详细文档

模型详情

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

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

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

训练详情

该模型从 经过语言建模适配的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许可证。