# Locret-Q

The test script for evaluating Locret-Q on RULER is adapted from Star Attention (https://github.com/NVIDIA/Star-Attention). We deeply thank them for providing a great and extendable evaluation framework!

## Setup

- First, set up a conda environment by
  ```
  conda env create -f environment.yaml
  conda activate locretq
  ```

- Then, setup the model path in `launch_locret.sh:17`
  ```
  MODEL_PATH="yourpath"
  ```

## Running Experiment

We have setup all the experimental script in `launch_locret.sh`.
```
# performance
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a flash -pc $PROMPT_CONFIG -l 131072
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a minf -pc $PROMPT_CONFIG -l 131072
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a locret -pc $PROMPT_CONFIG -l 131072 # this is in locret-q mode
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a snapkv -pc $PROMPT_CONFIG -l 131072
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a h2o -pc $PROMPT_CONFIG -l 131072
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a sirllm -pc $PROMPT_CONFIG -l 131072

# speed
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a flash -pc $PROMPT_CONFIG -l 131072 --time
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a minf -pc $PROMPT_CONFIG -l 131072 --time
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a locret -pc $PROMPT_CONFIG -l 131072 --time
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a snapkv -pc $PROMPT_CONFIG -l 131072 --time
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a h2o -pc $PROMPT_CONFIG -l 131072 --time
python $SCRIPT -p $MODEL_PATH -n $MODEL_NAME -a sirllm -pc $PROMPT_CONFIG -l 131072 --time
```

Just execute a single command and you are all set!
```
bash launch_locret.sh
```
