Skip to content

axa-rev-research/QuestGrad

Repository files navigation

QuestGrad

This repository contains the reproducibility code for QuestGrad, the method described in QuestGrad: Agentic Adversarial Data Generation for LLMs.

We use QuestGrad as shorthand for adversarial question generation with TextGrad-style feedback updates: the question itself is optimized to maximize disagreement between a weak target model and a stronger document-grounded model.

The public repository is intentionally small. It keeps the LegalBench workflow from the paper and removes unrelated experiments, notebooks, historical logs, and working-directory artifacts.

What This Reproduces

The paper workflow has three stages:

  1. Generate contract-specific training data with QuestGrad, or generate the Independent Questions baseline.
  2. LoRA fine-tune a Llama 3.1 8B model on that generated JSON file.
  3. Evaluate the fine-tuned checkpoint on the LegalBench questions for the same contract.

The easiest way to run those stages is through the Gradio app:

python questgrad_app.py

The same stages are also available as command-line scripts.

Repository Contents

  • questgrad_app.py: Gradio app for generation, LoRA training, and evaluation.
  • opti_gradFINAL.py: QuestGrad data generation method used by the paper.
  • indep-ques.py: Independent Questions baseline generation.
  • run_training.py: LoRA fine-tuning wrapper.
  • eval_legalbench.py: LegalBench evaluation wrapper.
  • textgrad/: local TextGrad dependency used by QuestGrad generation.
  • data/: the three LegalBench contracts used in the paper, plus bundled 50-question QuestGrad and Independent Questions JSON files for each contract.

Generated checkpoints, logs, result files, caches, and credentials are ignored by git.

Requirements

Use a Linux machine with a CUDA GPU for fine-tuning and evaluation. The repository is set up for a repo-local uv environment and does not require a system Python installation.

You also need:

  • Azure/OpenAI-compatible credentials for litellm.
  • Access to meta-llama/Llama-3.1-8B on Hugging Face, or a local mirror path.
  • A weak-model endpoint for QuestGrad generation. The default app setting uses ollama-llama3.1:8b-text-fp16, which opti_gradFINAL.py converts to the LiteLLM model string ollama/llama3.1:8b-text-fp16.
  • A local LegalBench checkout at ../legalbench, because eval_legalbench.py reads prompts from ../legalbench/tasks/<task>/base_prompt.txt.

Setup

From the repository root:

uv python install --install-dir .uv-python "$(cat .python-version)"
uv venv --python ./.uv-python/cpython-$(cat .python-version)-linux-x86_64-gnu/bin/python .venv
source .venv/bin/activate
uv pip install -r requirements.txt

If the local uv Python tree is incomplete or came from a broken machine restore, reinstall it with:

uv python install --reinstall --install-dir .uv-python "$(cat .python-version)"

Create the local environment file:

cp .env-example .env

Then fill in the Azure/OpenAI-compatible values in .env. This file is ignored by git and should not be committed.

Set up LegalBench next to this repository:

cd ..
git clone https://github.com/HazyResearch/legalbench.git legalbench
cd QuestGrad-repro

GUI Tutorial

Start the app:

source .venv/bin/activate
python questgrad_app.py

Open the local Gradio URL printed in the terminal.

1. Generate Training Data

Use the Generate Data tab.

For QuestGrad:

  • Select one of the three included contracts.
  • Keep the default weak, strong, and feedback model names if your server matches the paper-style setup.
  • Use 5 optimization epochs and 50 questions for the bundled paper-style configuration.
  • Click Generate QuestGrad Data.

For the Independent Questions baseline:

  • Select the same contract.
  • Use the Azure/OpenAI-compatible model name for question and answer generation.
  • Use 50 questions for the paper-style baseline.
  • Click Generate Independent Questions.

Generated JSON files are written to data/. The app also includes bundled 50-question JSON files, so you can skip generation and start from training.

2. Train a LoRA Checkpoint

Use the Train LoRA tab.

  • Select a JSON dataset from data/.
  • Select meta-llama/Llama-3.1-8B, or enter an equivalent local model path in the command-line workflow.
  • Keep the default training values for the paper-style run: 10 epochs, batch size 1, LoRA rank 32, LoRA alpha 32, dropout 0.05, and BF16 enabled.
  • Optionally set an output directory. If left empty, run_training.py creates one under finetuned-models/.
  • Click Start LoRA Training.

The app lists completed model directories and checkpoints from finetuned-models/.

3. Evaluate the Checkpoint

Use the Evaluate tab.

  • Select the same contract used for generation and training.
  • Select the fine-tuned model directory.
  • Select the checkpoint to evaluate.
  • Keep cuda:0 unless you want a different GPU.
  • Click Evaluate Fine-Tuned Checkpoint.

Evaluation metrics are written to results/.

Command-Line Equivalents

Generate one QuestGrad question and append it to a JSON file:

python opti_gradFINAL.py \
  ollama-llama3.1:8b-text-fp16 \
  azure-gpt-4o-mini \
  azure-gpt-4o-mini \
  "data/CardlyticsInc_20180112_S-1_EX-10.16_11002987_EX-10.16_Maintenance Agreement1.pdf" \
  "Cardl_QuestGrad_llama3.1:8b-text-fp16_azure-gpt-4o-mini_azure-gpt-4o-mini_5epochs_50questions.json" \
  5

opti_gradFINAL.py appends one generated item per invocation. The app repeats that command until the requested number of questions has been generated.

Generate the Independent Questions baseline:

python indep-ques.py \
  --contract_path "CardlyticsInc_20180112_S-1_EX-10.16_11002987_EX-10.16_Maintenance Agreement1.pdf" \
  --model azure/gpt-4o-mini \
  --num_questions 50

Fine-tune on a generated dataset:

python run_training.py \
  "data/Cardl_QuestGrad_llama3.1:8b-text-fp16_azure-gpt-4o-mini_azure-gpt-4o-mini_5epochs_50questions.json" \
  --base_model meta-llama/Llama-3.1-8B \
  --num_epochs 10 \
  --batch_size 1 \
  --bf16 \
  --output_dir finetuned-models/cardl_qg

Evaluate a checkpoint:

python eval_legalbench.py \
  --base-model meta-llama/Llama-3.1-8B \
  --checkpoint finetuned-models/cardl_qg/checkpoint-<STEP> \
  --contract "CardlyticsInc_20180112_S-1_EX-10.16_11002987_EX-10.16_Maintenance Agreement1.pdf" \
  --batch-size 8 \
  --device cuda:0 \
  --output-file results/cardl_eval.json

Notes On Reproducibility

The bundled JSON files let you reproduce the training and evaluation path without calling hosted generation models. Regenerating data can produce small differences because hosted LLMs, Ollama model builds, and sampling behavior may change over time.

For a paper-style run, use the non-instruct Llama 3.1 8B base model unless you are intentionally testing the instruct variant.

About

No description, website, or topics provided.

Resources

License

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages