## Setup
```
pip install -e .
pip install lm_eval sentencepiece protobuf
```

To obtain the accuracies/perplexities of existing models (with the sensitivity records in `sensitivities.pt` and `weight-extraction/sens/<model_name>`), directly go to Step 3.

## Step 1. Getting column-wise sensitivities (requires GPUs)
Go to `experiments/sensitivity-extraction`
Open extract_sensitivities and change the models variable to the list of models you want to extract the sensitivity of
Run `extract_sensitivities.sh`
It will create `sensitivities.pt`. This file contains column sensitivities (created by torch.save). If you run extract_sensitivities.sh the second time, it will just append new results if the model is new, it will not override them. If the model is the same, it will override the sensitivities

## Step 2. Getting layer-wise sensitivities 
Go to `experiments/weights-extraction`
Open `extract_and_aggregate_eff_multiple.sh`
Again, edit the list of models that you want to get the sensitivities of
Now open `extract_and_aggregate_eff.sh`
Change the number of samples as you wish, also change csv-path to an appropriate name for the model. It is recommended to have one file per model family (llama, gemma). However, it’s okay to have one file and multiple model sizes
Run `extract_and_aggregate_eff_multiple.sh`
Now open `convert_sensitivities.py` and make sure that the right csv is read
Run `convert_sensitivities.py`
Now you have layer sensitivities saved to `weight-extraction/sens/<model_name>/enable_oa_True.csv`

## Step 3. Getting accuracies/perplexities 
Go to `experiments/k-bits-reproduction`
First check `experiment_config.py` . There you will find the models to run, the tasks and the numerical formats configurations.
Choose the models and the tasks you want to test on (don’t choose both perplexity and accuracy tasks, either only perplexity or accuracy. This is for historical reasons
Choose the configurations that you want to run. A configuration is just a dict with a bunch of parameters, the configurations are stored in the `NUMERICAL_FORMATS_CONFIG` array
You can notice that I generate configuration for different bitwidth. Unfortunately, there is no direct way in the config for now to set the bitwidth, you need to calculate the ratios manually to get the bitwidth that you want. You have to do this only if you decide to add a model. If you need to do that, check draw.ipynb , I left the instructions there how to do that, it’s not that hard
Open `run.sh` . There are two ways to run the code
In an infinite loop, with all the logging and stuff. It’s in the loop because sometimes there is an OOM, and then the runai machine gets killed for being idle.
For debugging purposes, just run the command. It doesn’t use the cache, so it will run everything from scratch.
Run `run.sh` . The results are stored in `intermediate/results.pt` . With the flag `--use-cache` it will check if a model and a configuration is already there, and in this case it will skip the computation, so you can simply re-run the script if it crashes. However, without this flag it will override your result, so be careful.
To draw the results, go to `draw.ipynb`.