如何用BridgeTower将纯文本转为2048维向量用于LanceDB检索

要让BridgeTower输出纯文本的2048维向量，核心思路是给模型传入占位图像来满足处理器的输入要求，同时保证文本分支正常工作，输出的向量和图文联合生成的向量处于同一特征空间，不影响检索效果。

具体实现步骤 #

• 当处理纯文本时，生成符合模型输入尺寸的空白图像（比如384x384，BridgeTower常用输入尺寸）
• 将空白图像和文本一起传入处理器，构造模型输入
• 正常调用模型获取pooler_output，得到2048维向量

修改后的代码 #

from PIL import Image
import torch
from transformers import AutoProcessor, BridgeTowerModel

def bt_embedding_from_local_pretrained(prompt, image_path=None):
    model_name = "D:\\download\\bridgetower-large-itm-mlm-itc"
    processor = AutoProcessor.from_pretrained(model_name)
    model = BridgeTowerModel.from_pretrained(model_name)
    
    if image_path is not None and image_path != '':
        image = Image.open(image_path)
    else:
        # 生成空白占位图像，尺寸匹配模型要求（可根据模型实际要求调整）
        image = Image.new('RGB', (384, 384), color=(0, 0, 0))
    
    inputs = processor(images=image, text=prompt, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = outputs.pooler_output
    embeddings_list = embeddings.squeeze().tolist()
    
    return embeddings_list

# 测试纯文本生成向量
text_query = "A young boy playing on a playground with chairs and a bench"
text_vector = bt_embedding_from_local_pretrained(text_query)
print(len(text_vector))  # 输出2048

原理说明 #

BridgeTower的多模态架构中，文本分支和图像分支并行处理后融合，即使传入占位图像，模型的文本编码器仍会正常提取文本特征，最终的pooler_output包含文本的有效特征，且维度和图文联合输入的结果一致，完全可以用于LanceDB的向量检索。

内容的提问来源于stack exchange，提问作者王大爷