bge-vl 模型在线测试

目的

基于 bge-vl-large 和 bge-vl-base 模型在线对图片进行检索，提供可视化页面进行操作

功能

1. 支持从页面上传图片，存储到服务器上，同时对图片进行embedding 后存储到 faiss 中用于之后的检索
2. 支持在页面输入文本、文本+图像 进行检索，并支持用户选择具体的模型，默认模型为 bge-vl-large，检索后展示前topN 最相似的照片
3. 可以在页面上查看所有已经上传的图片

配置

1. 模型的目录，图片的目录、embedding文件的存储均可以通过配置文件指定

技术选型

使用 python的 streamlit 来写页面，左侧sidebar有查询和图片浏览即可

安装与使用

安装依赖

pip install -r requirements.txt

运行应用

有两种方式运行应用：

1. 完整版应用（包含图像检索功能）：

streamlit run app.py

2. 简化版应用（仅包含图像上传和浏览功能，不依赖BGE-VL模型）：

streamlit run simple_app.py

应用将在本地启动，通常在 http://localhost:8501 访问。

使用说明

1. 上传图片：在"Upload"页面上传图片，点击"Process Uploads"按钮处理上传的图片
2. 搜索图片：在"Search"页面可以通过文本、图片或两者结合进行搜索
3. 浏览图片：在"Browse"页面可以查看所有已上传的图片
4. 选择模型：在侧边栏可以选择使用的模型（bge-vl-large 或 bge-vl-base）

项目结构

• app.py: 主应用程序（完整版）
• simple_app.py: 简化版应用程序（不依赖BGE-VL模型）
• model_utils.py: 模型操作工具类
• config.py: 配置文件
• test_model.py: 模型测试脚本
• data/: 数据目录
  • uploads/: 上传的图片
  • embeddings/: 存储的嵌入向量和索引
  • models/: 缓存的模型文件

常见问题

模型加载错误

如果遇到模型加载错误，可能是由于以下原因：

1. 依赖冲突：torch、torchvision 和 transformers 版本不兼容
2. 内存不足：BGE-VL-large 模型较大，可能需要较多内存
3. Python 版本兼容性：建议使用 Python 3.8-3.10

解决方案：

• 使用 simple_app.py 进行图片上传和浏览
• 尝试使用较小的 BGE-VL-base 模型
• 运行 test_model.py 脚本测试模型加载

参考demo

import torch
from transformers import AutoModel

MODEL_NAME = "BAAI/BGE-VL-base" # or "BAAI/BGE-VL-large"

model = AutoModel.from_pretrained(MODEL_NAME, trust_remote_code=True) # You must set trust_remote_code=True
model.set_processor(MODEL_NAME)
model.eval()

with torch.no_grad():
    query = model.encode(
        images = "./assets/cir_query.png", 
        text = "Make the background dark, as if the camera has taken the photo at night"
    )

    candidates = model.encode(
        images = ["./assets/cir_candi_1.png", "./assets/cir_candi_2.png"]
    )
    
    scores = query @ candidates.T
print(scores)