处理文本列

PyTorch Frame 通过使用文本嵌入模型来处理文本列，这些模型可以是预训练的语言模型。 我们支持两种主要的文本嵌入模型使用选项：

1. 在物化阶段将文本预编码为嵌入（以便在训练阶段模型参数被冻结）

2. 在训练阶段生成文本嵌入并微调其模型参数。

这些选项有各自的权衡。选项（1）允许更快的训练，而选项（2）允许更准确的预测，但由于需要对文本模型进行微调，训练成本更高。在PyTorch Frame中，我们可以通过简单地指定其stype来为每个文本列指定使用哪个选项：在传递给Dataset的col_to_stype参数中，我们可以为希望使用选项（1）的列指定stype.text_embedded，为使用选项（2）的列指定stype.text_tokenized。让我们使用一个真实世界的数据集来学习如何实现这一点。

处理真实世界数据集中的文本列

PyTorch Frame 提供了一系列带有文本列的表格基准数据集，例如 MultimodalTextBenchmark。

正如我们简要讨论的那样，PyTorch Frame 为文本列提供了两种语义类型：

1. stype.text_embedded 将在数据集物化阶段使用用户指定的文本嵌入模型对文本进行预编码。

2. stype.text_tokenized 将在数据集物化阶段使用用户指定的文本分词器对文本进行分词。分词后的文本（整数序列）在训练阶段被输入到文本模型中，并对文本模型的参数进行微调。

初始化和实现数据集的过程类似于示例介绍。下面我们强调每种语义类型的差异。

将文本预编码为嵌入

对于stype.text_embedded，首先你需要指定文本嵌入模型。 这里，我们使用SentenceTransformer包。

pip install -U sentence-transformers

指定文本嵌入器

接下来，我们创建一个文本编码器类，该类以迷你批次的方式将字符串列表编码为文本嵌入。

from typing import List
import torch
from torch import Tensor
from sentence_transformers import SentenceTransformer

class TextToEmbedding:
    def __init__(self, device: torch.device):
        self.model = SentenceTransformer('all-distilroberta-v1', device=device)

    def __call__(self, sentences: List[str]) -> Tensor:
        # Encode a list of batch_size sentences into a PyTorch Tensor of
        # size [batch_size, emb_dim]
        embeddings = self.model.encode(
            sentences,
            convert_to_numpy=False,
            convert_to_tensor=True,
        )
        return embeddings.cpu()

然后我们实例化TextEmbedderConfig，它指定了我们用来预编码文本的text_embedder和batch_size。

from torch_frame.config.text_embedder import TextEmbedderConfig

device = (torch.device('cuda' if torch.cuda.is_available() else 'cpu')

col_to_text_embedder_cfg = TextEmbedderConfig(
    text_embedder=TextToEmbedding(device),
    batch_size=8,
)

请注意，基于Transformer的文本嵌入模型通常对GPU内存需求较高， 因此指定一个合理的batch_size（例如8）是很重要的。 此外，请注意我们将默认在所有文本列中使用相同的TextEmbedderConfig。 如果我们想为不同的文本列使用不同的text_embedder （例如"text_col0"和"text_col1"），我们可以 使用如下字典：

# Prepare text_embedder0 and text_embedder1 for text_col0 and text_col1, respectively.
col_to_text_embedder_cfg = {
    "text_col0":
    TextEmbedderConfig(text_embedder=text_embedder0, batch_size=4),
    "text_col1":
    TextEmbedderConfig(text_embedder=text_embedder1, batch_size=8),
}

为数据集嵌入文本列

一旦指定了col_to_text_embedder_cfg，我们可以将其传递给 Dataset对象，如下所示。

import torch_frame
from torch_frame.datasets import MultimodalTextBenchmark

dataset = MultimodalTextBenchmark(
    root='/tmp/multimodal_text_benchmark/wine_reviews',
    name='wine_reviews',
    col_to_text_embedder_cfg=col_to_text_embedder_cfg,
)

dataset.feat_cols  # This dataset contains one text column `description`
>>> ['description', 'country', 'province', 'points', 'price']

dataset.col_to_stype['description']
>>> <stype.text_embedded: 'text_embedded'>

然后我们调用dataset.materialize(path=...)，它将使用文本嵌入模型根据给定的col_to_text_embedder_cfg预编码text_embedded列。

# Pre-encode text columns based on col_to_text_embedder_cfg. This may take a while.
dataset.materialize(path='/tmp/multimodal_text_benchmark/wine_reviews/data.pt')

len(dataset)
>>> 105154

# Text embeddings are stored as MultiNestedTensor
dataset.tensor_frame.feat_dict[torch_frame.embedding]
>>> MultiNestedTensor(num_rows=105154, num_cols=1, device='cpu')

强烈建议在materialize()期间指定path。 它将缓存生成的TensorFrame，从而避免在每次运行materialize时嵌入文本，这可能会非常耗时。 一旦缓存，TensorFrame可以用于后续的materialize()调用。

注意

在内部，text_embedded 被分组到父类型 embedding 中，位于 TensorFrame 内。

将文本嵌入融合到表格学习中

PyTorch Frame 提供了 LinearEmbeddingEncoder 用于在 TensorFrame 中编码 embedding。该模块对预计算的嵌入应用线性函数。

from torch_frame.nn.encoder import (
    EmbeddingEncoder,
    LinearEmbeddingEncoder,
    LinearEncoder,
)

stype_encoder_dict = {
    stype.categorical: EmbeddingEncoder(),
    stype.numerical: LinearEncoder(),
    stype.embedding: LinearEmbeddingEncoder()
}

然后，stype_encoder_dict 可以直接输入到 StypeWiseFeatureEncoder。

微调文本模型

与stype.text_embedded相比， stype.text_tokenized在数据集物化阶段只进行最小限度的处理， 即仅对原始文本进行分词，将字符串转换为整数序列。 然后，在训练阶段，成熟的文本模型将分词后的句子作为输入， 并输出文本嵌入，这使得文本模型能够以端到端的方式进行训练。

在这里，我们使用 Transformers 包。

pip install transformers

指定文本分词

在stype.text_tokenized中，文本列将在数据集物化阶段进行分词。 让我们首先创建一个分词类，该类将字符串列表分词为torch.Tensor的字典。

from typing import List
from transformers import AutoTokenizer
from torch_frame.typing import TextTokenizationOutputs

class TextToEmbeddingTokenization:
    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')

    def __call__(self, sentences: List[str]) -> TextTokenizationOutputs:
        # Tokenize batches of sentences
        return self.tokenizer(
            sentences,
            truncation=True,
            padding=True,
            return_tensors='pt',
        )

在这里，输出 TextTokenizationOutputs 是一个字典， 其中键包括 input_ids 和 attention_mask，值 包含 PyTorch 的标记和注意力掩码张量。

然后我们为我们的文本嵌入模型实例化 TextTokenizerConfig 如下所示。

from torch_frame.config.text_tokenizer import TextTokenizerConfig

col_to_text_tokenizer_cfg = TextTokenizerConfig(
    text_tokenizer=TextToEmbeddingTokenization(),
    batch_size=10_000,
)

这里 text_tokenizer 将句子列表映射为 torch.Tensor 的字典， 这些张量在训练时作为文本模型的输入。 分词处理是在小批量中进行的，其中 batch_size 表示批量大小。 由于文本分词器在 CPU 上运行速度快，我们可以在这里指定相对较大的 batch_size。 另外，请注意，我们允许为不同的文本列指定 text_tokenizer 的字典， 使用 stype.text_tokenized。

# Prepare text_tokenizer0 and text_tokenizer1 for text_col0 and text_col1, respectively.
col_to_text_tokenizer_cfg = {
    "text_col0":
    TextTokenizerConfig(text_tokenizer=text_tokenizer0, batch_size=10000),
    "text_col1":
    TextTokenizerConfig(text_tokenizer=text_tokenizer1, batch_size=20000),
}

为数据集中的文本列进行分词

一旦指定了 col_to_text_tokenizer_cfg，我们可以将其传递给 Dataset 对象，如下所示。

import torch_frame
from torch_frame.datasets import MultimodalTextBenchmark

dataset = MultimodalTextBenchmark(
    root='/tmp/multimodal_text_benchmark/wine_reviews',
    name='wine_reviews',
    text_stype=torch_frame.text_tokenized,
    col_to_text_tokenizer_cfg=col_to_text_tokenizer_cfg,
)

dataset.col_to_stype['description']
>>> <stype.text_tokenized: 'text_tokenized'>

然后我们调用dataset.materialize()，它将使用文本分词器根据给定的col_to_text_tokenizer_cfg对text_tokenized列进行预分词。

# Pre-encode text columns based on col_to_text_tokenizer_cfg.
dataset.materialize()

# A dictionary of text tokenization results
dataset.tensor_frame.feat_dict[torch_frame.text_tokenized]
>>> {'input_ids': MultiNestedTensor(num_rows=105154, num_cols=1, device='cpu'), 'attention_mask': MultiNestedTensor(num_rows=105154, num_cols=1, device='cpu')}

请注意，我们使用了一个MultiNestedTensor的字典来存储分词结果。 我们使用字典的原因是，常见的文本分词器通常会返回多个文本模型输入，例如之前展示的input_ids和attention_mask。

使用表格学习微调文本模型

PyTorch Frame 提供了 LinearModelEncoder，旨在灵活地以每列方式应用任何可学习的 PyTorch 模块。我们首先指定一个 ModelConfig 对象，该对象声明了要应用于每列的模块。

注意

ModelConfig 有两个参数需要指定： 首先，model 是一个可学习的 PyTorch 模块，它接受 TensorFrame 中的每列张量作为输入， 并输出每列的嵌入。形式上，model 接受一个形状为 [batch_size, 1, *] 的 TensorData 对象作为输入， 然后输出形状为 [batch_size, 1, out_channels] 的嵌入。然后，out_channels 指定了 model 的输出嵌入维度。

我们可以使用上述的LinearModelEncoder功能来嵌入stype.text_tokenized在TensorFrame中。

要使用该功能，我们首先需要为ModelConfig准备model。 这里我们使用PEFT包和 LoRA策略来微调底层文本模型。

pip install peft

然后我们将model设计为一个使用DistilBERT和LoRA微调的模型。 请注意，model需要将每列的feat作为输入，并输出大小为[batch_size, 1, out_channels]的嵌入。 正如我们提到的，每列的feat在stype.text_tokenized的情况下是MultiNestedTensor的字典格式。 在forward()期间，我们首先通过使用to_dense()将每个MultiNestedTensor转换为填充的torch.Tensor，填充值由fill_value指定。

import torch
from torch import Tensor
from transformers import AutoModel
from torch_frame.data import MultiNestedTensor
from peft import LoraConfig, TaskType, get_peft_model

class TextToEmbeddingFinetune(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = AutoModel.from_pretrained('distilbert-base-uncased')
        # Set LoRA config
        peft_config = LoraConfig(
            task_type=TaskType.FEATURE_EXTRACTION,
            r=32,
            lora_alpha=32,
            inference_mode=False,
            lora_dropout=0.1,
            bias="none",
            target_modules=["ffn.lin1"],
        )
        # Update the model with LoRA config
        self.model = get_peft_model(self.model, peft_config)

    def forward(self, feat: dict[str, MultiNestedTensor]) -> Tensor:
        # Pad [batch_size, 1, *] into [batch_size, 1, batch_max_seq_len], then,
        # squeeze to [batch_size, batch_max_seq_len].
        input_ids = feat["input_ids"].to_dense(fill_value=0).squeeze(dim=1)
        # Set attention_mask of padding idx to be False
        mask = feat["attention_mask"].to_dense(fill_value=0).squeeze(dim=1)

        # Get text embeddings for each text tokenized column
        # out.last_hidden_state has the shape:
        # [batch_size, batch_max_seq_len, out_channels]
        out = self.model(input_ids=input_ids, attention_mask=mask)

        # Use the CLS embedding to represent the sentence embedding
        # Return value has the shape [batch_size, 1, out_channels]
        return out.last_hidden_state[:, 0, :].unsqueeze(1)

现在我们已经准备好了model。我们可以通过额外提供out_channels参数来实例化ModelConfig对象。在DistilBERT的情况下，out_channels是768。

from torch_frame.config import ModelConfig
model_cfg = ModelConfig(model=TextToEmbeddingFinetune(), out_channels=768)

然后我们指定col_to_model_cfg，将每个列名映射到所需的model_cfg。

col_to_model_cfg = {"description": model_cfg}

我们现在可以将col_to_model_cfg传递给LinearModelEncoder，以便它将指定的model应用于所需的列。在这种情况下，我们将模型TextToEmbeddingFinetune应用于stype.text_tokenized列，该列在TensorFrame中被称为"description"。

from torch_frame.nn import (
    EmbeddingEncoder,
    LinearEncoder,
    LinearModelEncoder,
)

stype_encoder_dict = {
    stype.categorical: EmbeddingEncoder(),
    stype.numerical: LinearEncoder(),
    stype.text_tokenized: LinearModelEncoder(col_to_model_cfg=col_to_model_cfg),
}

生成的 stype_encoder_dict 可以直接输入到 StypeWiseFeatureEncoder 中。

请参考 pytorch-frame/examples/transformers_text.py 以获取更多关于使用Transformers包进行文本嵌入和微调的信息。

此外，请参考 pytorch-frame/examples/llm_embedding.py 以获取更多关于大型语言模型的文本嵌入信息，例如 OpenAI embeddings 和 Cohere embed。