# RotatE-Box


## Install

```bash
cd KBI-Regex
pip install -e .
```


## Experiments

To run the experiments, pre-training on single hop queries is required. The default parameters set are for FB15K.

To train RotatE-box on FB15K single hop queries

```bash
python main.py --model RotatE --box --dataset fb15k -dim 800 --query_types 0 --do_train -adv --lr_schedule half --seed 0 --save_dir <YOUR_OUTPUT_DIR>
```

There is discrepency in `dim` in the code base and in the paper, `-dim 800` is equivalent to `k=400` in the paper.

The checkpoint obtained from single hop queries will be used in regex query training.

To train RotatE-box(COMP) on FB15K regex queries:

```bash
python main.py --model RotatE --box --kleene_plus_op gqe --disjunction_op gqe --dataset fb15k -dim 800 --query_types 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 --do_train --resume_from_checkpoint <PATH_TO_KBC_CHECKPOINT> --seed 0 --save_dir <YOUR_OUTPUT_DIR>
```

Options for `--kleene_plus_op` are `gqe`, `free_param`. Options for `--disjunction_op` are `gqe`, `aggregation`. `gqe` corresponds to both Projection for Kleene plus and Deepsets for disjunction.

Apart from the setting above (`gqe`, `gqe`), for any other setting, remove `20` from ```query_types`.

To run RotatE models, remove `--box` from above commands. To run Query2Box, replace `RotatE` with `Query2Box`.

In the codebase, Wiki100 is referred as `wiki_v2`. The complete dataset couldn't be included due to size restrictions. Only dev and test sets are available.
