You can deploy this model and start earning money today!

RAG-Gym/
ReSearch-MedQA-SFT

Name: ReSearch-MedQA-SFT
Availability: InStock

peft

safetensors

1910.09700

meta-llama/Llama-3.1-8B-Instruct

adapter

Views

No views yet

ReSearch-MedQA-SFT

This is the SFT actor for the Re$^2$Search agent on MedQA, released as part of RAG-Gym. Re$^2$Search is the proposed architecture in RAG-Gym. It is the Re$^2$Search (Reason, Reflection, Search) agent, which reflects on its reasoning to surface unsupported claims and turn them into targeted queries. This checkpoint was trained with supervised fine-tuning on preferred process-reward actions.

Base model: meta-llama/Meta-Llama-3.1-8B-Instruct
Agent architecture: Re$^2$Search (the proposed architecture in RAG-Gym)
Task / dataset: MedQA (medical multiple-choice question answering in the USMLE style)
Training method: supervised fine-tuning (SFT)
Adapter: PEFT LoRA (r=256, alpha=512), task type CAUSAL_LM
Precision: bf16

Intended use

Act as the reasoning/search policy for the Re$^2$Search agent on MedQA-style (medical multiple-choice question answering in the USMLE style) tasks within RAG-Gym.
Can be run with zero-shot inference or, together with the matching PRM critic, with critic-guided (Best-of-N) inference.

This model is intended for research on process-supervised retrieval-augmented generation. It is not intended for clinical decision-making or other high-stakes use.

How to use

Load the LoRA adapter on top of the base model:

python

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

base = "meta-llama/Meta-Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(base)
model = AutoModelForCausalLM.from_pretrained(base, torch_dtype=torch.bfloat16)
model = PeftModel.from_pretrained(model, "RAG-Gym/ReSearch-MedQA-SFT")

For full usage within the agent framework (environment setup, inference, and critic-guided Best-of-N selection), see the RAG-Gym repository.

Training

Fine-tuned from meta-llama/Meta-Llama-3.1-8B-Instruct on MedQA process-reward data using supervised fine-tuning (SFT) with the TRL library. Adaptation uses LoRA (rank 256, alpha 512, task type CAUSAL_LM) in bf16. The training code and full configuration are available in the RAG-Gym repository.

Limitations

The model is trained on MedQA-style data and is intended for research use within RAG-Gym. Outputs may be inaccurate or unsupported by retrieved evidence and should not be relied upon for high-stakes decisions. As an 8B-parameter model, it requires a CUDA-enabled GPU to run efficiently.

License

Public Domain Notice (U.S. Government Work, NCBI). See the LICENSE in the repository. Please cite the authors when using this material.

Citation

bibtex

@article{xiong2025raggym,
    title={RAG-Gym: Optimizing Reasoning and Search Agents with Process Supervision},
    author={Guangzhi Xiong and Qiao Jin and Xiao Wang and Yin Fang and Haolin Liu and Yifan Yang and Fangyuan Chen and Zhixing Song and Dengyu Wang and Minjia Zhang and Zhiyong Lu and Aidong Zhang},
    journal={arXiv preprint arXiv:2502.13957},
    year={2025}
}