RobustRAG / scripts /VisRAG_code.md
AaronCIH's picture
Upload folder using huggingface_hub
a8c50ce verified

跑通

安裝VisRAG環境

安裝VisRAG.

git clone https://github.com/OpenBMB/VisRAG.git
conda create --name VisRAG python==3.10.8
conda activate VisRAG
conda install nvidia/label/cuda-11.8.0::cuda-toolkit
cd VisRAG
pip install -r requirements.txt
pip install -e .
cd timm_modified
pip install -e .
cd ..
  1. 把原先的 "/VisRAG/src", "/VisRAG/script" 更名或刪除, 換成我寫的 "/mnt/191/a/lyw/visrag2/VisRAG/src" , "/mnt/191/a/lyw/visrag2/VisRAG/scripts"

  2. pip install peft==0.10.0 # lora (印象太高版本無法跑, 0.10.0 還行)

  3. pip install scikit-image==0.25.2 # degradation 會用到

  4. pip install opencv-python==4.11.0.86 # degradation 會用到

  5. 改掉 training script 跟 eval script 的路徑

    a. 改掉 "/VisRAG/scripts/train_retriever/train.sh" 裡面的 "/mnt/VENV/user-conda/r12943109/miniconda3/envs/loramoe/lib" 跟 "/mnt/VENV/user-conda/r12943109/miniconda3/envs/loramoe/bin/python" 換成你conda的路徑.(共兩處)
    b. 改掉 "/VisRAG/scripts/eval_retriever"的: eval.sh, eval2.sh, evalreal.sh, evalreal2.sh, evalreal3.sh  裡面的 "/mnt/VENV/user-conda/r12943109/miniconda3/envs/loramoe/bin/python"換成你的 python 路徑.(共兩處) 
  1. 去 "/VisRAG/src/openmatch/arguments.py" 改 "*cache_dir" 成自己的路徑 (可用ctrl+f 尋找 /mnt/191/ 有兩個)

  2. 下載 我們的model(VisIR)(LLM+vpm1+vpm2+zero_linear)

    a. VisIR(乾淨未訓練, zero_linear全0): "/mnt/191/a/lyw/VisRAG/openbmb/VisIR" (整個資料夾都要)
    b. VisIR(stage1): "/mnt/191/a/lyw/VisRAG/checkpoints/train-2025-06-19-065530-model-data-lr-1e-5-softm_temp-0.02-bsz16-ngpus4-nnodes1-inbatch--nepoch-1-pooling-wmean-attention-causal-qinstruct-true-cinstruct-false-gradcache-true-passage-stopgrad-false-npassage-1" (整個資料夾都要)
    c. VisIR(stage1+2):  "/mnt/191/a/lyw/VisRAG/checkpoints/train-2025-06-19-194437-model-data-lr-1e-5-softm_temp-0.02-bsz16-ngpus4-nnodes1-inbatch--nepoch-1-pooling-wmean-attention-causal-qinstruct-true-cinstruct-false-gradcache-true-passage-stopgrad-false-npassage-1" (整個資料夾都要)

eval指令

eval腳本的差異如下

eval script 有五個, 分別對應5個測試資料集
1. eval.sh: 可測乾淨的 ArxivQA,ChartQA,MP-DocVQA,InfoVQA,PlotQA,SlideVQA
2. eval2.sh: 可測degraded的 ArxivQA,ChartQA,MP-DocVQA,InfoVQA,PlotQA,SlideVQA
3. evalreal.sh: 可測乾淨的 rvl-cdip 
4. evalreal2.sh: 可測realworld的 rvl-cdip 
5. evalreal3.sh: 可測realworld的 MP-DocVQA
但其實他們的差別只有在CHECKPOINT_DIR(儲存測試結果位置), CORPUS_PATH, QUERY_PATH, QRELS_PATH 而已

eval的使用方法

PYTHONNOUSERSITE=1 bash <測試的腳本路徑> 512 2048 <每台gpu測多少batchsize> <使用幾個gpu> wmean causal <測試的資料集名稱> <儲存時的子資料夾名稱(辨識用,否則都是日期)> <model dir路徑> <goat lora safetensors路徑,沒有就空著,注意這個要指到 .safetensors 不是dir>

eval.sh跟eval2.sh的使用方法

<測試的資料集名稱> 最多可以指定6個資料集 "ArxivQA,ChartQA,MP-DocVQA,InfoVQA,PlotQA,SlideVQA" 以逗號隔開

e.g. 
PYTHONNOUSERSITE=1 bash scripts/eval_retriever/eval.sh 512 2048 16 2 wmean causal ArxivQA,ChartQA,MP-DocVQA,InfoVQA,PlotQA,SlideVQA 乾淨visir測試  "/mnt/191/a/lyw/VisRAG/openbmb/VisIR"

e.g. 
PYTHONNOUSERSITE=1 bash scripts/eval_retriever/eval2.sh 512 2048 16 2 wmean causal ArxivQA,ChartQA,MP-DocVQA,InfoVQA,PlotQA,SlideVQA 乾淨visir測試  "/mnt/191/a/lyw/VisRAG/openbmb/VisIR"

e.g. 
PYTHONNOUSERSITE=1 bash scripts/eval_retriever/eval2.sh 512 2048 16 2 wmean causal ArxivQA,ChartQA,MP-DocVQA,InfoVQA,PlotQA,SlideVQA 乾淨visir測試  "/mnt/191/a/lyw/VisRAG/openbmb/VisIR" "/goat/lora/path/如果有的話.safetensors"

evalreal.sh跟evalreal2.sh跟evalreal3.sh的使用方法

<測試的資料集名稱> 不重要, 但至少需指定一個, 腳本不會使用這個input但會被用在取名上, 但建議是用 "ChartQA" 這個, 你打六個的話她會重複測六次

e.g. 
PYTHONNOUSERSITE=1 bash scripts/eval_retriever/evalreal.sh 512 2048 16 2 wmean causal ChartQA 乾淨visir測試  "/mnt/191/a/lyw/VisRAG/openbmb/VisIR"  
e.g. 
PYTHONNOUSERSITE=1 bash scripts/eval_retriever/evalreal2.sh 512 2048 16 2 wmean causal ChartQA 乾淨visir測試  "/mnt/191/a/lyw/VisRAG/openbmb/VisIR"  
e.g. 
PYTHONNOUSERSITE=1 bash scripts/eval_retriever/evalreal3.sh 512 2048 16 2 wmean causal ChartQA 乾淨visir測試  "/mnt/191/a/lyw/VisRAG/openbmb/VisIR"  

eval的輔助腳本(如果忘記紀錄的話, 只要有.trec檔案就可以重新計算)

可使用 "/src/把visrag_eval_trec_算出三個分.py",
假設: 
/VisRAG/evaloutput/乾淨visir測試/eval-...-abc/ArxivQA
/VisRAG/evaloutput/乾淨visir測試/eval-...-abc/ChartQA
/VisRAG/evaloutput/乾淨visir測試/eval-...-abc/MP-DocVQA
/VisRAG/evaloutput/乾淨visir測試/eval-...-abc/InfoVQA
/VisRAG/evaloutput/乾淨visir測試/eval-...-abc/PlotQA
/VisRAG/evaloutput/乾淨visir測試/eval-...-abc/SlideVQA

只要改腳本中的 這三個就可以
ROOT=/VisRAG/evaloutput/乾淨visir測試/eval-...-abc/
CACHE_DIR = "/path/to/save/your/datasets/"
要看什麼 = 'recall_10' or 要看什麼 = 'recip_rank'

=> 就可以重現結果

train指令

train 指令跟 eval 指令不同就一個腳本而已. "/VisRAG/scripts/train_retriever/train.sh"

train.sh 裡面唯一需要手動改的只有 --save_steps 500. 你需要設定 你要多少steps保存一個 checkpoints
訓練完後他會保存在 /VisRAG/checkpoints/... 裡面

train 指令大致使用方法

"有些不需要動的我就不標示了"
PYTHONNOUSERSITE=1  bash scripts/train_retriever/train.sh 2048 \
    <每個gpu的batchsize> <gpu數量> 0.02 1 true false config/deepspeed.json \
    1e-5 false wmean causal 1 true <minibatch> false \
    <乾淨的模型dir> <資料集名稱> \
    <是否要使用degradation> <是否要訓練小的資料集就好> <現在是stage1還是stage2> <有沒有使用lora>

2048是passage使用的token數量
0.02是訓練時的溫度, 仿visrag
1e-5是learning-rate, 仿visrag
<資料集名稱>: 就兩個 "openbmb/VisRAG-Ret-Train-In-domain-data" or "openbmb/VisRAG-Ret-Train-Synthetic-data"
<minibatch>: 可設1,2,3,4,..., 意思是要inference batchsize 個 sample 其實是由 minibatch 累積達到的. 設越大越耗vram
<是否要使用degradation>: true or false, 設true就好, 就是說在finetune的時候要不要把資料變成degraded的
<是否要訓練小的資料集就好>: true or false, 設false就好, 設true的話她只會train前30k筆, 可以去"/VisRAG/src/openmatch/dataset/train_dataset.py"改30k成其他數字
<現在是stage1還是stage2>: "stage1" or "stage2": 用來標示現在是在訓練 "stage1" 還是 "stage2" 注意不能打錯
<有沒有使用lora>: "None" or "任意的非None字段": 用來標示有沒有用到lora

stage1 訓練

e.g. 
PYTHONNOUSERSITE=1  bash scripts/train_retriever/train.sh 2048 \
    16 2 0.02 1 true false config/deepspeed.json \
    1e-5 false wmean causal 1 true 1 false \
    '/mnt/191/a/lyw/VisRAG/openbmb/VisIR' 'openbmb/VisRAG-Ret-Train-In-domain-data' \
    true false 'stage1' 'None'

stage2 訓練

e.g. 
PYTHONNOUSERSITE=1  bash scripts/train_retriever/train.sh 2048 \
    16 2 0.02 1 true false config/deepspeed.json \
    1e-5 false wmean causal 1 true 1 false \
    '/path/to/stage1/訓練完的model/dir' 'openbmb/VisRAG-Ret-Train-In-domain-data' \
    true false 'stage2' '有用lora'

程式碼部分

GOAT CONFIG 的部分:

注意我的GOAT 的 config 用幾個 experts 那些參數 是在 code 裡面做 setting 的
可以參考"/VisRAG/src/openmatch/modeling/dense_retrieval_model.py": build()
裡面有一行, goat_config = GoatV3Config(), 所以基本上目前要自己記錄下來 config 是怎麼調整
或者你可以去改"/VisRAG/src/openmatch/modeling/goat_injector.py" 裡面的 GoatV3Config 的 default 參數

訓練時的程式碼:

1. "/VisRAG/src/openmatch/driver/train.py": 主要的訓練腳本
2. "/VisRAG/src/openmatch/arguments.py": 想確認arguments有哪些可以去那邊看
3. "/VisRAG/src/openmatch/modeling/modeling_visir/modeling_vis_ir.py": 我模型的定義
4. "/VisRAG/src/openmatch/modeling/goat_injector.py": goat moe lora的實現
5. "/VisRAG/src/openmatch/dataset/train_dataset.py":line213(def get_process_fn): 那邊定義我是怎處理huggingface訓練資料集 
6. "/VisRAG/src/openmatch/modeling/dense_retrieval_model.py": build(): 在這裡我定義我是怎麼拿到要訓練的model的
    注意, 他是先用DRModel把 retrieval model 給包起來,  所以在程式碼中看到 model(...)的時候, 其實是由 DRModel的forward處理
    可參考 DRModel的 forward() & encode()
7. "/VisRAG/src/openmatch/trainer/dense_trainer.py": _save(...): 在這邊定義我是怎麼儲存model的
8. "/VisRAG/src/openmatch/trainer/dense_trainer.py": training_step(...): 在這邊定義訓練model的過程
    因為採用 minibatch 的方式, 他會先全部inference一次拿到cache, 再重新inference計算loss

eval時的程式碼:

1. "/VisRAG/src/openmatch/driver/eval.py": 主要的推理腳本
2. "/VisRAG/src/openmatch/arguments.py": 想確認arguments有哪些可以去那邊看
    eval 的時候沒用到 train_args, 可以用 if  train_args: 來區分是否是訓練
3. "/VisRAG/src/openmatch/modeling/modeling_visir/modeling_vis_ir.py": 我模型的定義
4. "/VisRAG/src/openmatch/modeling/goat_injector.py": goat moe lora的實現
5. "/VisRAG/src/openmatch/dataset/inference_dataset.py": 在這邊跟inference處理資料有關
    注意: 這邊我新增一個, "特殊處理用" 的 variable,  就是把 realworld datasets 全部轉正向這樣
6. "/VisRAG/src/openmatch/retriever/dense_retriever.py": 在這邊定義了他是怎麼先把embedding存下來, 最後再算分

模型架構圖

乾淨沒加可參考 "/mnt/191/a/lyw/visrag2/VisIR(模型架構).txt"
加了GOAT可參考 "/mnt/191/a/lyw/visrag2/VisIR+GOAT(模型架構).txt"
想調整GOAT架構可以改
"/VisRAG/src/openmatch/modeling/goat_injector.py" 裡面的 GoatV3Config 或者
"/VisRAG/src/openmatch/modeling/dense_retrieval_model.py" 直接在 build(...) 的程式碼裡面更改加載邏輯