T4

Tell As You Want: Customizing Image Narrative with Knowledge and Thoughts [paper] [supp]

Pretrained weights

You can download the pretrained weights from the Huggingface: T4.

Dataset

You can download the dataset from the Huggingface: ImgNarr-23K.

sample：
{
    "data_id": "Unique ID",
    "question": "User question",
    "answer": "Standard answer",
    "image_path": "Unzip 'images.zip' to a localpath and replace 'localpath/' with 'xx/'",
    "query_psg_gt": "Reference documents used during training",
    "gemini_cot": "Contains <CAPTION><SUB-QA><CONCLUSION>",
    "narrative": "Contains <CONCLUSION>"
}

For the definitions of other fields, please refer to InfoSeek and EVQA.

Knowledge Bases and Index Construction

We provide the knowledge bases used in T4. We employ a two-stage retrieval approach that first conducts coarse document retrieval and then performs fine-grained passage retrieval.

Stage 1: coarse document retrieval

In this stage, we retrieve documents relevant to the image using EVA-CLIP model.

InfoSeek : Following ReflectiVA, we utilize the 100K knowledge base, which is a filtered subset of the original 6M for knowledge base.

Encyclopedic-VQA : We use original 2M knowledge base.

You can find the pre-built index files in the ReflectiVA.

Stage 2: fine-grained passage retrieval

In this stage, we further retrieve the passages most relevant to the question from the documents related to the image using PreFLMR model.

InfoSeek : We segment the 100K documents from the first stage into 200-word chunks, resulting in the file Wiki100K_ver_1_0_200.json.

Encyclopedic-VQA : We utilize the paragraphs from the original 2M knowledge base, the pre-processed paragraphs are available in M2KR.

Index Construction

Please follow the installation instructions in PreFMLR.

Index a custom document collection:

cd src_pub/search/FLMR
python examples/index_docs.py

Installation

# [Optional] Create a clean Conda environment
conda create -n t4 python==3.10.0
conda activate t4

# Install PyTorch and dependencies (make sure CUDA version matches)
pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu126

# Install remaining dependencies
pip install -r t4_requirements.txt

Inference

Step 1: Start retriever

bash scripts/start_retriever.sh

Set the path in start_flmr_server_main.py to your local path.

Step 2: Start vllm

bash scripts/start_vllm.sh

Set MODEL_NAME_OR_PATH in start_vllm.sh to your local inference model path.

Step 3: Start inference

Standard inference using EVA-PreFLMR retrieval:

bash scripts/start_infer.sh

Fast inference using pre-retrieved content:

bash scripts/start_infer_context.sh

Set the input data path in src/data_loader.py to your local path.

Finetuning

cd train
pip install llama-recipes
bash ft.sh

Remember to modify the data_paths and image_base_path in datasets/cot_dataset.py to your local training datase path.

Evaluation

pip install git+https://github.com/openai/CLIP.git
pip install git+https://github.com/jmhessel/pycocoevalcap.git
pip install tensorflow
cd eval
python eval.py

Acknowledgements

Thanks to the code of CoRAG, LLaVA-CoT, FLMR and data of ReflectiVA, Encyclopedic-VQA and InfoSeek.