region:us
# Representation Classifier for tool selection
 
This repository contains code for training and evaluating a Representation Classifier (RepC) for the task of tool selection. The RepC is trained to classify user queries into different tool categories based on the descriptions of tools provided.

## Usage

### Training
To train the Representation Classifier, you can use the following command:

```bash
python repc.py --train --model_path teknium/OpenHermes-2.5-Mistral-7B --task tool --data_path data/data_within_tool_subset.json --n_train 5 --classifier svm --model_save_dir saved_models --use_cache
```

- `model_path`: Path to the pretrained language model.
- `task`: Which task to run.
- `data_path`: Path to the dataset containing tool descriptions and queries.
- `n_train`: Number of training examples for **each class**.
- `classifier`: Which classifier to use, only support svm currently.
- `model_save_dir`: Directory to save trained models.
- `use_cache`: Whether to use cached embeddings.

### Evaluation
To evaluate the trained Representation Classifier, you can use the following command:

```bash
python repc.py --evaluate --model_path teknium/OpenHermes-2.5-Mistral-7B --task tool --data_path data/data_within_tool_subset.json --n_train 5 --classifier svm --model_save_dir saved_models --use_cache
```

- `model_path`: Path to the pretrained language model.
- `task`: Which task to run.
- `data_path`: Path to the dataset containing tool descriptions and queries.
- `n_train`: Number of training examples for **each class**.
- `classifier`: Which classifier to use, only support svm currently.
- `model_save_dir`: Directory to save trained models.
- `use_cache`: Whether to use cached embeddings.

### Zero-shot Baseline
To run the zero-shot baseline, use the following command:

```bash
python repc.py --zero_baseline --model_path teknium/OpenHermes-2.5-Mistral-7B --task tool --data_path data/data_within_tool_subset.json --n_train 5
```

- `model_path`: Path to the pretrained language model.
- `task`: Which task to run.
- `data_path`: Path to the dataset containing tool descriptions and queries.
- `n_train`: Number of training examples for **each class**. This is only used for test set split.

### Few-shot Baseline
To run the few-shot baseline, use the following command:

```bash
python repc.py --few_shot_baseline --model_path teknium/OpenHermes-2.5-Mistral-7B --task tool --data_path data/data_within_tool_subset.json --n_train 5 --num_examples 3
```

- `model_path`: Path to the pretrained language model.
- `task`: Which task to run.
- `data_path`: Path to the dataset containing tool descriptions and queries.
- `n_train`: Number of training examples for **each class**. This is only used for test set split.
- `num_examples`: Number of examples as demonstration.

### Train and evaluate
You can also train and evaluate RepC and baselines in one command:

```bash
python repc.py --train --evaluate --zero_baseline --few_shot_baseline --model_path teknium/OpenHermes-2.5-Mistral-7B --task tool --data_path data/data_within_tool_subset.json --n_train 5 --num_examples 3 --classifier svm --model_save_dir saved_models --use_cache
```

## Inference
We also provided an inference code (inference.py) for end-to-end prediction:
```python
from repc import *

tools = ['CatIndexTool', 'SearchAlertsTool', 'VisualizationTool', 'SearchAnomalyDetectorsTool', 'SearchAnomalyResultsTool', 'SearchMonitorsTool', 'PPLTool', 'RAGTool']
questions = ["What is the number of documents in the index .kibana_1?", "How many alerts have severity level 1?"]
repc = RepC(model_path="teknium/OpenHermes-2.5-Mistral-7B", task="tool", device="auto")
predictions = repc.predict(classifier="svm", model_path="saved_models_n5/svm_l13.pkl", input=questions)
predictions = [tools[p] for p in predictions]
print(predictions)
```

## Results
The evaluation results are shown below, the original log can be found in tool_selection_results.txt.
```
| Method                | Accuracy (%) | Macro F1 (%) |
|-----------------------|--------------|--------------|
| RepC-layer            | 94.99        | 89.04        |
| RepC-ensemble         | 91.02        | 81.94        |
| Zero-baseline         | 10.86        | 25.95        |
| 3-shot-baseline       | 6.47         | 31.06        |
```