ImageTextRetrieval

提供商ohgnues
分类image-text-retrieval
许可证Apache-2.0
下载量5
星标0

简介

ImageTextRetrieval 是一款专注于图文跨模态检索的轻量化模型，旨在解决“用词找图”或“以图搜词”的实际需求。它不同于生成式 AI（如 Midjourney），而是通过将图像和文本映射到同一个向量空间，实现高效的语义匹配。对于开发者而言，该模型非常适合集成到私有图库搜索、电商商品检索或多模态知识库中。由于采用了 Apache-2.0 开源协议，企业部署成本低，且上手难度较小，可直接作为向量检索系统的底层编码器使用。

核心亮点

  • 实现精准的图文互搜，支持语义级匹配
  • Apache-2.0 协议，商业部署无压力
  • 适用于构建私有化多模态检索系统
  • 低延迟响应，可替代传统的关键词标签搜索

使用方法

安装依赖
# 安装 Hugging Face transformers
pip install transformers torch
SDK 使用
# 使用 transformers 加载模型
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("ohgnues/ImageTextRetrieval")
tokenizer = AutoTokenizer.from_pretrained("ohgnues/ImageTextRetrieval")

Hugging Face 下载

我们推荐使用命令行或者 Hugging Face Hub SDK 来进行模型的下载。

操作指引：在下载前，请先通过如下命令安装 huggingface_hub：

操作指引 
pip install -U huggingface_hub

命令行下载

下载完整模型库

下载完整模型库 
huggingface-cli download ohgnues/ImageTextRetrieval

下载单个文件到指定本地文件夹（以下载 config.json 到当前路径下 ./dir 目录为例）

下载单个文件到指定本地文件夹（以下载 config.json 到当前路径下 ./dir 目录为例） 
huggingface-cli download ohgnues/ImageTextRetrieval config.json --local-dir ./dir

更多命令行下载选项，可参见官方文档

SDK 下载

SDK 下载 
# 模型下载
from huggingface_hub import snapshot_download
model_dir = snapshot_download('ohgnues/ImageTextRetrieval')

Git 下载

请确保 lfs 已经被正确安装

Git 下载 
git lfs install
git clone https://huggingface.co/ohgnues/ImageTextRetrieval

如果您希望跳过 lfs 大文件下载，可以使用如下命令

跳过 LFS 
GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/ohgnues/ImageTextRetrieval

模型文件托管在 Hugging Face Hub，使用 HF CLI / SDK / Git 直接下载，不经过本站。

PyTorch / Transformers 使用

安装 Transformers

安装 Transformers 
pip install -U transformers torch

模型加载和推理

模型加载和推理 
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained('ohgnues/ImageTextRetrieval')
tokenizer = AutoTokenizer.from_pretrained('ohgnues/ImageTextRetrieval')

完整文档

来源: HuggingFace

Repo

multi-modal-retrieval

This repository contains code for multi modal retrieval

This project involves implementing a multi-modal Bi-encoder using both ResNet and BERT for image and text representations.

Data

Sample Data

The pretraining was conducted using the dataset from Hugging Face's "poloclub/diffusiondb" dataset.

I used 50k randomly sampled images and prompts for my project.

If you want to use a different dataset, follow the steps below

Data Format

Only images and the corresponding text for those images are necessary, and other elements are irrelevant. In this case, the text can serve as prompts or captions for the images.

You specify the names of the columns for images and text in the training command.

bash
python3 train.py --text_column_name text --image_column_name img

Pretrained models

Pretrained models can be downloaded huggingface or Specify the model name "ohgnues/ImageTextRetrieval" in the training command.

bash
python3 train.py --pretrained_model_name_or_path ohgnues/ImageTextRetrieval

The model "ohgnues/ImageTextRetrieval" was trained for 10 epochs using a Tesla P100 GPU.

Usage

Train

bash
python3 train.py --name 2m_random_50k --cache_dir /data/.cache --max_length 100 --num_train_epochs 10
For detailed instructions, please refer to the official Hugging Face documentation or consult the dataclass within the "train.py" script.

Encode

python
def encode(self, model_name: Literal["text", "image"],
            input_ids: Optional[torch.Tensor] = None,
            attention_mask: Optional[torch.Tensor] = None,
            token_type_ids: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.Tensor] = None,
            head_mask: Optional[torch.Tensor] = None,
            inputs_embeds: Optional[torch.Tensor] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
            pixel_values: Tensor = None
            ):
        
        if model_name == "text":
            return self.text_encoder(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            ).last_hidden_state[:, 0, :]
        
        elif model_name == "image":
            return self.image_encoder(
            pixel_values=pixel_values,
            output_hidden_states=output_hidden_states,
            ).pooler_output[:, :, 0, 0]
查看官方来源