## A typical code retrieval pipeline
Setting up a code retrieval pipeline requires multiple steps:

1. Download the document library.
2. Embed each document and place it in a vector store.
3. Define a retriever which, given a query, embeds it and then looks up relevant documents in the vector store.
4. (Optional) Define a reranker which takes top N retrieved documents and reorders them.

## Use [Sage](https://github.com/Storia-AI/sage) to bootstrap your code retrieval pipeline
While you are free to implement the entire pipeline from scratch, you can use [Sage](https://github.com/Storia-AI/sage) for an off-the-shelf solution.

### Installation
To install Sage, create a virtual environment and run `pip install`:
```bash
python -m venv my-venv
source activate my-venv/bin/activate
pip install git+https://github.com/Storia-AI/sage.git@main
```

### Prerequisites
Follow the [Prerequisites](https://github.com/Storia-AI/sage?tab=readme-ov-file#prerequisites) section to set up your environment (including, for instance, the desired API keys).

### Indexing huggingface/transformers
To clone + embed + store the repository in a vector store, run this command:
```bash
sage-index huggingface/transformers --commit-hash=7bb1c99800d235791dace10305731f377db8077b
```
Note that the dataset was created using the commit above. If you choose to clone the `huggingface/transformers` library manually, make sure to use the same commit.

We defaulted `sage-index` to run locally on your machine. For higher quality, you might want to use state-of-the-art third-party providers like OpenAI for embeddings, etc. You can overwrite these defaults from the command line. For instance:

```bash
sage-index huggingface/transformers \
--commit-hash=7bb1c99800d235791dace10305731f377db8077b
--mode=remote \
--embeddings-provider=openai \
--embeddings-model=text-embedding-3-small
```

Please refer directly to Sage's documentation for more details. You're also welcome to fork the repository to add your own extensions (and we accept pull requests if you're particularly proud of your contributions!).

### Optional: chat with the codebase as a sanity-check

If you want to gain an intuition of how well your retrieval algorithm is working, you can bring up a chat interface by running the following command. This will bring up a Gradio app:

```bash
sage-chat huggingface/transformers --commit-hash=7bb1c99800d235791dace10305731f377db8077b
```

### Making a submission file
You can easily produce a submission file in the right format with the [retrieve_kaggle.py](https://github.com/Storia-AI/sage/blob/main/benchmarks/retrieval/retrieve_kaggle.py) script.