/xmc.dspy

In-Context Learning for eXtreme Multi-Label Classification (XMC) using only a handful of examples.

Primary LanguagePythonMIT LicenseMIT

Infer-Retrieve-Rank: In-Context Learning for Extreme Multi-Label classification

Resources:

  1. [Jan'23] Paper announcement thread & Code announcement thread
  2. [Jan'23] In-Context Learning for Extreme Multi-Label classification paper

Infer-Retrieve-Rank (IReRa) is a generic and modular program which specifies interactions between pretrained Language Models and Retrievers to efficiently solve multi-label classification tasks with an extreme amount of classes (≥10,000 classes). Using only ≅ 50 labeled input examples, IReRa can be optimized to achieve state-of-the-art performance, despite not being finetuned. This optimization generally involves having a strong, teacher Language Model (e.g. gpt4) solve the task and gather instructions or demonstrations that help a more efficient, student Language Model (e.g. llama-2) solve the task better. A user can easily specify which parts of the program are implemented using which LMs, to strike the perfect balance between cost and performance.

The goal of IReRa is to easily be applicable to a broad range of tasks involving inference with Language Models and Retrieval. To this end, this repository detangles (i) the logic of Infer-Retrieve-Rank, (ii) the prompts needed to adapt IReRa to a specific domain and (iii) optimization techniques to improve performance. All of this is made possible through the dspy programming model.

LM calls of our results are cached, meaning you can reproduce our runs without paying for any inference costs yourself.

Table of Contents

  1. Installation
  2. Load data
  3. Compile and run
  4. Apply to new tasks
  5. Writing custom programs
  6. Contributing
  7. Citation

1) Installation

Create the conda environment:

conda create -n xmc python=3.10
conda activate xmc

Install DSPy. Our work relies on an experimental branch of DSPy.

git clone -b irera --single-branch https://github.com/stanfordnlp/dspy.git
cd dspy/
git checkout 802f2d5f26c1a64d8aad6adbd8b4394b9c4bb743
pip install .
cd ..

Install the rest of the requirements:

pip install -r requirements.txt

(optional) Set up models if you plan to do new runs

All LM calls to reproduce our runs are cached. However, if you want to execute new calls you will need to set up OpenAI and local models.

Set your OpenAI API KEY:

export OPENAI_API_KEY=<your-key>

For local models (e.g. llama-2), set up a local Text Generation Interface (TGI) via these steps. Make sure to point the url field in lm_config.json to your TGI server.

2) Load data

Load the data and cache needed to reproduce our results.

bash scripts/load_data.sh
bash scripts/load_cache.sh

3) Compile and run

All compilation IReRa runs from our paper can be reproduced by running bash scripts/compile_left_to_right.sh or bash scripts/run_left_to_right.sh. We provide the resulting program state in results_precompiled/ so you can load our compiled IReRa programs.

Load our compiled IReRa for ESCO_TECH and evaluate:

python run_irera.py \
    --dataset_name esco_tech \
    --state_path ./results_precompiled/esco_tech_infer-retrieve-rank_00/program_state.json \
    --lm_config_path ./lm_config.json \
    --do_validation \
    --do_test

Compile your own IReRa on ESCO_TECH and evaluate:

python compile_irera.py \
    --dataset_name esco_tech \
    --ontology_name esco \
    --prior_path ./data/esco/esco_priors.json \
    --ontology_path ./data/esco/skills_en_label.txt \
    --infer_signature_name infer_esco \
    --rank_signature_name rank_esco \
    --retriever_model_name sentence-transformers/all-mpnet-base-v2 \
    --infer_student_model_name llama-2-13b-chat \
    --infer_teacher_model_name gpt-3.5-turbo-instruct \
    --rank_student_model_name gpt-4-1106-preview \
    --rank_teacher_model_name gpt-4-1106-preview \
    --infer_compile_metric_name rp10 \
    --rank_compile_metric_name rp10 \
    --prior_A 0 \
    --rank_topk 50 \
    --do_validation \
    --do_test \
    --optimizer_name left-to-right \
    --lm_config_path ./lm_config.json

Command line arguments are explained in the respective files.

If you want to speed up the runs, you can use multithreading (warning: this can mess up caching sometimes).

export DSP_NUM_THREADS=8

The results from run_irera.py are slightly different than those of compile_irera.py, most likely due to a minor bug in loading and saving models. We take the results of compile_irera.py as the official results we report in the paper.

4) Apply to new tasks

To apply IReRa to a new task, you minimally need to add a new dataset and write a custom signature

A) Adding a new dataset

Data loading logic is defined in src/data_loaders. Every loader should minimally return a validation and testing dataframe, where each row has a 'text' field (which is a string) and a 'label' field (which is a list of strings).

compile_irera expects a .txt file where each line is a label, via the --ontology_path argument. These labels can not contain commas, since this currently breaks embedding the labels in a comma-separated list. Make sure labels have no trailing spaces. Via --prior_path a dictionary needs to be supplied which maps labels to probabilities.

B) Writing custom signatures

A signature tells an in-context learning module what its task is and what the inputs and outputs look like. You can think of it as a minimal zero-shot prompt for the task. All signatures are defined in src/programs/signatures.py, add your new signature here.

To add a new signature, simple subclass dspy.Signature with your custom fields and description. Add your new class to the supported_signatures dictionary at the bottom of src/programs/signatures.py so it can be accessed throughout the code.

C) Adding new metrics

Raw metrics are defined in src/metrics.py and are wrapped with DSPy logic in src/evaluators.py. When adding a metric, make sure to add it to the supported_metrics dictionary at the bottom of src/evaluators.py so it can be used throughout the code.

5) Writing custom programs

The goal of IReRa is to be modular and easily applicable to new tasks. To alter IReRA, you minimally need to write custom logic for (new) modules write custom optimizers which bootstrap your program.

A) Writing custom logic

Infer-Retrieve-Rank is defined in src/programs/infer_retrieve_rank.py. It is initialized with a config, defined in src/programs/config.py. You can write new programs or manipulate the behavior of IReRa by changing the python code in the forward method.

When you introduce new hyperparameters, make sure to add these to the config you are using. Otherwise these will not be saved and loaded.

B) Writing custom optimizers

Optimization strategies are defined in src/optimizer.py. When adding a new optimizer, make sure to add it to the supported_optimizers dictionary at the bottom of the file, so it can be accessed throughout the code.

Currently, optimizers are aware of which modules the to-optimize program has. This needs to be resolved so optimization strategies can be flexibly applied to different programs.

6) Contributing

I am happy to take on collaborations! Feel free to contact me at the email defined in the paper, or by opening an issue or pull request on the GitHub page. Here are some issues we could work on:

Research:

  • Incorporate label definitions, synonyms, and hierarchy to improve performance
  • Applying IReRa as recommender systems (The Extreme Classification Repository has many such tasks)
  • Add ensembling strategies to IReRa to improve performance (both ensembling programs and ensembling different modules within those programs)
  • Add chunking / summarization strategies to deal with long inputs efficiently
  • Optimize different parts of the program for different metrics to improve end-to-end performance (e.g. optimize Infer-Retrieve for high Recall@50 and use Rank to improve Recall@10)
  • Figure out how to use best use labels for program output to bootstrap supervision on intermediate program steps for better optimization
  • Automatically pick the best LMs / retrievers from HuggingFace to implement a given module (AutoIReRa)

Software Engineering:

  • Add a HugginFace dataloader to apply IReRa to any multi-label classification task
  • Add a cli argument parsing class
  • Create a config for the Optimizer class
  • Make the Optimizer class agnostic of the implementation details of the program that gets optimized, so it can define a high-level strategy which gets compiled across many programs
  • Track amount of student and teacher calls during optimization so system-cost can be compared.
  • Make src/programs/retrieve.py more efficient
  • Track intermediate pipeline steps and log these traces for debugging
  • Use logging instead of print-statements, and dump logs to experiment file
  • Add RAG(atouille)
  • Control seeds in LM calls and data loaders

7) Citation

Get notified of future work by following @KarelDoostrlnck on Twitter.

If you found this repo useful or use Infer-Retrieve-Rank in your work, please cite us:

@article{d2024context,
  title={In-Context Learning for Extreme Multi-Label Classification},
  author={D'Oosterlinck, Karel and Khattab, Omar and Remy, Fran{\c{c}}ois and Demeester, Thomas and Develder, Chris and Potts, Christopher},
  journal={arXiv preprint arXiv:2401.12178},
  year={2024}
}