Update README and complete the inference implementation

Author: SkyFishMoon

.gitignore

@@ -176,3 +176,5 @@ checkpoints/
 wandb/
 
 dataset/
+generation_results/
+backups/

README.md

@@ -1,180 +1,44 @@
-# EasyR1: An Efficient, Scalable, Multi-Modality RL Training Framework
+# RAVQA
 
 [![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/EasyR1)](https://github.com/hiyouga/EasyR1/stargazers)
 [![Twitter](https://img.shields.io/twitter/follow/llamafactory_ai)](https://twitter.com/llamafactory_ai)
 
-This project is a clean fork of the original [veRL](https://github.com/volcengine/verl) project to support vision language models, we thank all the authors for providing such a high-performance RL training framework.
+## Document VQA
 
-EasyR1 is efficient and scalable due to the design of **[HybirdEngine](https://arxiv.org/abs/2409.19256)** and the latest release of **[vLLM](https://github.com/vllm-project/vllm)**'s SPMD mode.
+### Dataset Preprocessing
 
-## Features
+#### Corpus Building
 
-- Supported models
-  - Llama3/Qwen2/Qwen2.5 language models
-  - Qwen2/Qwen2.5-VL vision language models
-  - DeepSeek-R1 distill models
+Change the raw data path and the target path in `rag_serving/build_corpus.py`
 
-- Supported algorithms
-  - GRPO
-  - Reinforce++
-  - ReMax
-  - RLOO
-
-- Supported datasets
-  - Any text, vision-text dataset in a [specific format](#custom-dataset)
-
-- Supported tricks
-  - Padding-free training
-  - Resuming from checkpoint
-  - Wandb & SwanLab & Mlflow & Tensorboard tracking
-
-## Requirements
-
-### Software Requirements
-
-- Python 3.9+
-- transformers>=4.51.0
-- flash-attn>=2.4.3
-- vllm>=0.8.3
-
-We provide a [Dockerfile](./Dockerfile) to easily build environments.
-
-We recommend using the [pre-built docker image](https://hub.docker.com/r/hiyouga/verl) in EasyR1.
-
-```bash
-docker pull hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.3-flashinfer0.2.2-cxx11abi0
-```
-
-### Hardware Requirements
-
-\* *estimated*
-
-| Method                   | Bits |  1.5B  |   3B   |   7B   |   32B   |
-| ------------------------ | ---- | ------ | ------ | ------ | ------- |
-| GRPO Full Fine-Tuning    |  AMP | 2*24GB | 4*40GB | 8*40GB | 16*80GB |
-| GRPO Full Fine-Tuning    | BF16 | 1*24GB | 1*40GB | 4*40GB |  8*80GB |
-
-> [!NOTE]
-> Use `worker.actor.fsdp.torch_dtype=bf16` and `worker.actor.optim.strategy=adamw_bf16` to enable bf16 training.
->
-> We are working hard to reduce the VRAM in RL training, LoRA support will be integrated in next updates.
-
-## Tutorial: Run Qwen2.5-VL GRPO on [Geometry3K](https://huggingface.co/datasets/hiyouga/geometry3k) Dataset in Just 3 Steps
-
-![image](assets/qwen2_5_vl_7b_geo.png)
-
-### Installation
-
-```bash
-git clone https://github.com/hiyouga/EasyR1.git
-cd EasyR1
-pip install -e .
-```
-
-### GRPO Training
-
-```bash
-bash examples/qwen2_5_vl_7b_geo3k_grpo.sh
-```
-
-### Merge Checkpoint in Hugging Face Format
-
-```bash
-python3 scripts/model_merger.py --local_dir checkpoints/easy_r1/exp_name/global_step_1/actor
-```
-
-> [!TIP]
-> If you encounter issues with connecting to Hugging Face, consider using `export HF_ENDPOINT=https://hf-mirror.com`.
->
-> If you want to use SwanLab logger, consider using `bash examples/qwen2_5_vl_7b_geo3k_swanlab.sh`.
-
-## Custom Dataset
-
-Please refer to the example datasets to prepare your own dataset.
-
-- Text dataset: https://huggingface.co/datasets/hiyouga/math12k
-- Image-text dataset: https://huggingface.co/datasets/hiyouga/geometry3k
-- Multi-image-text dataset: https://huggingface.co/datasets/hiyouga/journeybench-multi-image-vqa
-
-## How to Understand GRPO in EasyR1
-
-![image](assets/easyr1_grpo.png)
-
-- To learn about the GRPO algorithm, you can refer to [Hugging Face's blog](https://huggingface.co/docs/trl/v0.16.1/en/grpo_trainer).
-
-## How to Run 70B+ Model in Multi-node Environment
-
-1. Start the Ray head node.
-
-```bash
-ray start --head --port=6379 --dashboard-host=0.0.0.0
-```
-
-2. Start the Ray worker node and connect to the head node.
-
-```bash
-ray start --address=<head_node_ip>:6379
+```shell
+python rag_serving/build_corpus.py
 ```
 
-3. Check the Ray resource pool.
+#### Image Index Building
 
-```bash
-ray status
+```shell
+python index_builder.py --retrieval_method vdr-2b-v1 --model_path llamaindex/vdr-2b-v1 --corpus_path /scratch-scc/projects/scc_ulsb_fe/yang/images_corpus/images.parquet --save_dir /scratch-scc/projects/scc_ulsb_fe/yang/images_index --max_length 512 --batch_size 128 --faiss_type Flat --index_modal image --sentence_transformer  --save_embedding
 ```
 
-4. Run training script on the Ray head node only.
-
-```bash
-bash examples/qwen2_5_vl_7b_geo3k_grpo.sh
-```
-
-See the **[veRL's official doc](https://verl.readthedocs.io/en/latest/start/multinode.html)** for more details about multi-node training and Ray debugger.
-
-## Other Baselines
-
-We also reproduced the following two baselines of the [R1-V](https://github.com/deep-agent/R1-V) project.
-- [CLEVR-70k-Counting](examples/baselines/qwen2_5_vl_3b_clevr.sh): Train the Qwen2.5-VL-3B-Instruct model on counting problem.
-- [GeoQA-8k](examples/baselines/qwen2_5_vl_3b_geoqa8k.sh): Train the Qwen2.5-VL-3B-Instruct model on GeoQA problem.
-
-## Awesome Work using EasyR1
-
-- **MMR1**: Advancing the Frontiers of Multimodal Reasoning. [![[code]](https://img.shields.io/github/stars/LengSicong/MMR1)](https://github.com/LengSicong/MMR1)
-- **Vision-R1**: Incentivizing Reasoning Capability in Multimodal Large Language Models. [![[code]](https://img.shields.io/github/stars/Osilly/Vision-R1)](https://github.com/Osilly/Vision-R1) [![[arxiv]](https://img.shields.io/badge/arxiv-2503.06749-blue)](https://arxiv.org/abs/2503.06749)
-- **Seg-Zero**: Reasoning-Chain Guided Segmentation via Cognitive Reinforcement. [![[code]](https://img.shields.io/github/stars/dvlab-research/Seg-Zero)](https://github.com/dvlab-research/Seg-Zero) [![[arxiv]](https://img.shields.io/badge/arxiv-2503.06520-blue)](https://arxiv.org/abs/2503.06520)
-- **MetaSpatial**: Reinforcing 3D Spatial Reasoning in VLMs for the Metaverse. [![[code]](https://img.shields.io/github/stars/PzySeere/MetaSpatial)](https://github.com/PzySeere/MetaSpatial) [![[arxiv]](https://img.shields.io/badge/arxiv-2503.18470-blue)](https://arxiv.org/abs/2503.18470)
-- **Temporal-R1**: Envolving Temporal Reasoning Capability into LMMs via Temporal Consistent Reward. [![[code]](https://img.shields.io/github/stars/appletea233/Temporal-R1)](https://github.com/appletea233/Temporal-R1)
-- **NoisyRollout**: Reinforcing Visual Reasoning with Data Augmentation. [![[code]](https://img.shields.io/github/stars/John-AI-Lab/NoisyRollout)](https://github.com/John-AI-Lab/NoisyRollout) [![[arxiv]](https://img.shields.io/badge/arxiv-2504.13055-blue)](https://arxiv.org/pdf/2504.13055)
-- **GUI-R1**: A Generalist R1-Style Vision-Language Action Model For GUI Agents. [![[code]](https://img.shields.io/github/stars/ritzz-ai/GUI-R1)](https://github.com/ritzz-ai/GUI-R1) [![[arxiv]](https://img.shields.io/badge/arxiv-2504.10458-blue)](https://arxiv.org/abs/2504.10458)
-
-## TODO
-
-- Support LoRA (high priority).
-- Support ulysses parallelism for VLMs (middle priority).
-- Support more VLM architectures.
-
-> [!NOTE]
-> We will not provide scripts for supervised fine-tuning and inference in this project. If you have such requirements, we recommend using [LLaMA-Factory](https://github.com/hiyouga/LLaMA-Factory).
-
-### Known bugs
-
-These features are temporarily disabled for now, we plan to fix them one-by-one in the future updates.
-
-- Vision language models are not compatible with ulysses parallelism yet.
+### Launch RL
 
-## Discussion Group
+#### Tool Environment Serving
 
-👋 Join our [WeChat group](assets/wechat.jpg).
+1. Get the IP address of the server
 
-## FAQs
+    ```shell
+    hostname --ip-address
+    ```
 
-> ValueError: Image features and image tokens do not match: tokens: 8192, features 9800
+2. Start serving
 
-Increase the `data.max_prompt_length` or reduce the `data.max_pixels`.
+    ```shell
+    python rag_serving/serving.py --config rag_serving/serving_config.yaml --num_retriever 4 --port 42354
+    ```
 
-> RuntimeError: CUDA Error: out of memory at /workspace/csrc/cumem_allocator.cpp:62
+#### RL Training
 
-Reduce the `worker.rollout.gpu_memory_utilization` and enable `worker.actor.offload.offload_params`.
 
-> RuntimeError: 0 active drivers ([]). There should only be one.
 
-Uninstall `deepspeed` from the current python environment.
+## General VQA

vllm_rollout_backup.txt backups/vllm_rollout_backup.txtvllm_rollout_backup.txt renamed to backups/vllm_rollout_backup.txt

@@ -409,7 +409,7 @@ class vLLMRolloutAgent(vLLMRollout, ImageProcessMixin):
                         sampling_params=self.sampling_params,
                         use_tqdm=False
                     )
-                pydevd_pycharm.settrace('47.76.117.131', port=47508, stdoutToServer=True, stderrToServer=True)
+                # pydevd_pycharm.settrace('47.76.117.131', port=47508, stdoutToServer=True, stderrToServer=True)
                 search_queries = []
                 search_indices = []
                 search_doc_ids = []

vllm_rollout_backup3.txt backups/vllm_rollout_backup3.txtvllm_rollout_backup3.txt renamed to backups/vllm_rollout_backup3.txt

@@ -0,0 +1,116 @@
+#!/bin/bash
+#SBATCH --job-name=EasyR1-qwen2p5VL-7b-DocAgent
+#SBATCH --nodes=2
+#SBATCH --mem=450G
+#SBATCH --mail-user=tianyu.yang@uni-goettingen.de
+#SBATCH --mail-type=all
+#SBATCH --cpus-per-task=64
+#SBATCH -p kisski
+#SBATCH --gpus-per-node=4
+#SBATCH -t 48:00:00
+#SBATCH --output=slurm-%j.out
+#SBATCH --error=slurm-%j.err
+#############module load cuda/12.2.1
+############SBATCH --constraint=80gb
+################SBATCH --mem=500G
+
+set -x
+#export VLLM_ATTENTION_BACKEND=XFORMERS
+
+MODEL_PATH=checkpoints/EasyR1/global_step_355/actor/huggingface  # replace it with your local file path
+WANDB_API_KEY=a3b3f7b7962a8b549c4635ee3a03944d554f1a10
+ROLLOUT_NAME=vllm_agent
+SEARCH_TOP_N=1
+SEARCH_URL=http://10.241.148.102:42354
+LIMIT_IMAGES=15
+MAX_RESPONSE_LENGTH=15000
+MAX_PROMPT_LENGTH=1024
+ROLLOUT_MAX_NUM_BATCHED_TOKENS=16024
+TENSOR_PARALLEL_SIZE=2
+PROMPT_KEY=question
+ROLLOUT_BATCH_SIZE=128
+ROLLOUT_N=1
+VAL_BATCH_SIZE=-1
+TEMPERATURE=0.2
+TEST_DATA_PATH=/mnt/vast-kisski/projects/kisski-sub-doc-understanding/EasyR1/dataset/test/feta.parquet
+
+CONFIG_PATH=/mnt/vast-kisski/projects/kisski-sub-doc-understanding/EasyR1/examples/generation_config.yaml
+SAVE_PATH=/mnt/vast-kisski/projects/kisski-sub-doc-understanding/EasyR1/generation_results/qwen2_5_vl_7b_doc_agent_test
+
+if [ "$WANDB_API_KEY" != "None" ]; then
+    wandb login --relogin $WANDB_API_KEY
+fi
+
+# make output directory
+if [ ! -d "$SAVE_PATH" ]; then
+    mkdir -p $SAVE_PATH
+fi
+
+nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+nodes_array=($(scontrol show hostnames "$SLURM_JOB_NODELIST"))
+
+head_node=${nodes_array[0]}
+head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)
+
+if [[ "$head_node_ip" == *" "* ]]; then
+IFS=' ' read -ra ADDR <<<"$head_node_ip"
+if [[ ${#ADDR[0]} -gt 16 ]]; then
+  head_node_ip=${ADDR[1]}
+else
+  head_node_ip=${ADDR[0]}
+fi
+echo "IPV6 address detected. We split the IPV4 address as $head_node_ip"
+fi
+
+port=6379
+ip_head=$head_node_ip:$port
+export ip_head
+echo "IP Head: $ip_head"
+
+
+echo "StartingHEAD at $head_node"
+srun --nodes=1 --ntasks=1 -w "$head_node" /bin/bash -c \
+       "source /user/yang28/u14705/.bashrc && source /mnt/vast-kisski/projects/kisski-sub-doc-understanding/miniconda3/bin/activate EasyR1 \
+        && ray start --head --node-ip-address="$head_node_ip" --port=$port \
+         --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_NODE}" --include-dashboard true --dashboard-host 0.0.0.0 --dashboard-port 8265 --block" &
+# optional, though may be useful in certain versions of Ray < 1.0.
+sleep 10
+
+# number of nodes other than the head node
+worker_num=$((SLURM_JOB_NUM_NODES - 1))
+#export worker_num = 1
+
+for ((i = 1; i <= worker_num; i++)); do
+    node_i=${nodes_array[$i]}
+    echo "Starting WORKER $i at $node_i"
+    srun --nodes=1 --ntasks=1 -w "$node_i" /bin/bash -c \
+      "source /user/yang28/u14705/.bashrc && source /mnt/vast-kisski/projects/kisski-sub-doc-understanding/miniconda3/bin/activate EasyR1  \
+      && ray start --address "$ip_head" --num-cpus "${SLURM_CPUS_PER_TASK}" --num-gpus "${SLURM_GPUS_PER_NODE}" --block" &
+    sleep 5
+done
+
+
+srun --overlap --nodes=1 --ntasks=1 -w "$head_node"  /bin/bash -c \
+  "source /user/yang28/u14705/.bashrc && source /mnt/vast-kisski/projects/kisski-sub-doc-understanding/miniconda3/bin/activate EasyR1  \
+  && python -m verl.trainer.main_generation \
+    config=${CONFIG_PATH} \
+    data.test_files=${TEST_DATA_PATH} \
+    data.prompt_key=${PROMPT_KEY} \
+    data.format_prompt=./examples/format_prompt/doc_agent.py \
+    data.max_response_length=${MAX_RESPONSE_LENGTH} \
+    data.max_prompt_length=${MAX_PROMPT_LENGTH} \
+    data.rollout_batch_size=${ROLLOUT_BATCH_SIZE} \
+    worker.actor.model.model_path=${MODEL_PATH} \
+    worker.rollout.tensor_parallel_size=${TENSOR_PARALLEL_SIZE} \
+    worker.rollout.name=${ROLLOUT_NAME} \
+    worker.rollout.n=${ROLLOUT_N} \
+    worker.rollout.temperature=${TEMPERATURE} \
+    worker.rollout.max_num_batched_tokens=${ROLLOUT_MAX_NUM_BATCHED_TOKENS} \
+    worker.rollout.top_n=${SEARCH_TOP_N} \
+    worker.rollout.search_url=${SEARCH_URL} \
+    worker.rollout.limit_images=${LIMIT_IMAGES} \
+    worker.reward.score_function=./examples/score_function/doc_agent.py:compute_score \
+    trainer.n_gpus_per_node=${SLURM_GPUS_PER_NODE} \
+    trainer.nnodes=${SLURM_NNODES} \
+    trainer.save_checkpoint_path=${SAVE_PATH}"
+#    trainer.load_checkpoint_path=/mnt/vast-kisski/projects/kisski-sub-doc-understanding/EasyR1/checkpoints/qwen2_5_vl_7b_doc_agent/global_step_160"

vllm_rollout_backup_2.txt backups/vllm_rollout_backup_2.txtvllm_rollout_backup_2.txt renamed to backups/vllm_rollout_backup_2.txt

@@ -1 +1 @@
-system_prompt = "You are a helpful assistant designed to answer user questions based on a user-provided multi-page document. The document can not be input directly with the question, you must reason step by step to determine how to obtain evidence document pages by optimally utilizing tools and analyze the relevant content in the obtained document pages to precisely answer user's question. Your reasoning process MUST BE enclosed within <think> </think> tags. Your answer MUST BE enclosed within <answer> </answer> tags. In the last part of the answer, the final exact answer is enclosed within \\boxed{{}} with latex format. The available tools include a **search tool** and a **fetch tool**. During reasoning, you can invoke either the search tool by generating <search> your search query here </search> to retrieve document pages most relevant to your search query or or the fetch tool by generating <fetch> page number </fetch> to obtain a specific document page. For example, your response could be in the format of \'<think> your reasoning process </think> <search> search query </search>\', or \'<think> your reasoning process </think> <fetch> page number </fetch>\', or \'<think> your reasoning process </think> <answer> your answer here. The final answer is \\[ \\boxed{{answer here}} \\] </answer>\'. After invoking a tool, the user will return obtained document pages inside <result> </result> tags to you.\n\n**Important constraints**:\n- If there is no answer found in the document, respond with <answer> The final answer is \\[ \\boxed{{The problem is not answerable}} \\] </answer>.\n- If multiple valid answers are found, return them separated by semicolons.\n- Only one page can be fetched at a time using the fetch tool.\n- Enrich the user question to form a good search query to get more accurate retrieval results.\n- Do not naively use the fetch tool if you don't know the specific page number of the document page that the user is asking about."
+system_prompt = "You are a helpful assistant designed to answer user questions based on a user-provided multi-page document. The document can not be input directly with the question, you must reason step by step to determine how to obtain evidence document pages by optimally utilizing tools and analyze the relevant content in the obtained document pages to precisely answer user's question. Your reasoning process MUST BE enclosed within <think> </think> tags. Your answer MUST BE enclosed within <answer> </answer> tags. In the last part of the answer, the final exact answer is enclosed within \\boxed{{}} with latex format. The available tools include a **search tool** and a **fetch tool**. During reasoning, you can invoke either the search tool by generating <search> your search query here </search> to retrieve document pages most relevant to your search query or or the fetch tool by generating <fetch> page number </fetch> to obtain a specific document page. For example, your response could be in the format of \'<think> your reasoning process </think> <search> search query </search>\', or \'<think> your reasoning process </think> <fetch> page number </fetch>\', or \'<think> your reasoning process </think> <answer> your answer here. The final answer is \\[ \\boxed{{answer here}} \\] </answer>\'. After invoking a tool, the user will return obtained document pages inside <result> </result> tags to you.\n\n**Important constraints**:\n- If there is no answer found in the document, respond with <answer> The final answer is \\[ \\boxed{{The problem is not answerable}} \\] </answer>.\n- If multiple valid answers are found, return them separated by semicolons.\n- Only one page can be fetched at a time using the fetch tool.\n- Enrich the user question to form a good search query to get more accurate retrieval results.\n- Do not naively use the fetch tool if  you don't know the specific page number of the document page that the user is asking about."