We provide the code used to run OPRM inference.

The main script is eval_longbench_oprm.py, which evaluates OPRM on Longbench with Recurrent-Gemma-IT-9B.


#####
Setup
#####

1. All dependencies for the script are found in requirements.txt, which can be installed in a conda environment via:
conda env create -f environment.yml -p /path/to/conda/env

2. After installing the environment, please replace:
/path/to/conda/env/lib/python3.10/site-packages/transformers/generation/utils.py

With the provided utils.py (code added for OPRM is marked by <OPRM inference> ... <\OPRM inference>)


3. Also replace:
/path/to/conda/env/lib/python3.10/site-packages/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py

With the provided modeling_recurrent_gemma.py




#######
Running
#######

Arguments:
* --device - cuda device (e.g. 'cuda:0')
* --e - toggles Longbench \ Longbench_e (0 for LB, 1 for LB_e)
* --is_oprm - toggles oprm (0 for vanilla inference, 1 for OPRM inference) 


Running Example:
python eval_longbench_oprm.py --device 'cuda:0' --e 1 --is_oprm 1


