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許可證。