(End2End RAG) Finetuning Script (finetune_rag_ray_end2end.sh) not working for HF Dataset
Author: YichiRockyZhang
Pull request: no

### System Info

Python 3.9.16
Transformers 4.13.0
WSL

### Who can help?

@shamanez  @ArthurZucker

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Note: I have gotten all test scripts to work.

I'm running the below modified `finetune_rag_ray_end2end.sh` script. I specifically have the `wiki_dpr/psgs_w100.multiset.compressed` version of the `wiki_dpr` dataset cached (with train generated).

```bash
# Sample script to finetune RAG using Ray for distributed retrieval.

# Add parent directory to python path to access lightning_base.py
export PYTHONPATH="../":"${PYTHONPATH}"

#creates the custom knowlegebase
# python use_own_knowledge_dataset.py  \
#     --csv_path /DIR/SQUAD-KB/squad-kb.csv \
#     --output_dir  /DIR/SQUAD-KB

# Start a single-node Ray cluster.
ray start --head

python finetune_rag.py \
    --data_dir  ./ft_data \
    --output_dir ./model_checkpoints \
    --model_name_or_path facebook/rag-token-base \
    --model_type rag_token \
    --fp16 \
    --gpus 4  \
    --profile \
    --do_train \
    --end2end \
    --do_predict \
    --n_val -1  \
    --train_batch_size 32 \
    --eval_batch_size 32 \
    --max_source_length 300 \
    --max_target_length 25 \
    --val_max_target_length 25 \
    --test_max_target_length 25 \
    --label_smoothing 0.1 \
    --dropout 0.1 \
    --attention_dropout 0.1 \
    --weight_decay 0.001 \
    --adam_epsilon 1e-08 \
    --max_grad_norm 0.1 \
    --lr_scheduler polynomial \
    --learning_rate 3e-05 \
    --num_train_epochs 10 \
    --warmup_steps 500 \
    --gradient_accumulation_steps 8 \
    --distributed_retriever ray \
    --num_retrieval_workers 4  \
    --index_name hf \
    --index_path /path/to/.hf_cache/datasets/wiki_dpr/psgs_w100.multiset.compressed/0.0.0/74d4bff38a7c18a9498fafef864a8ba7129e27cb8d71b22f5e14d84cb17edd54/psgs_w100.multiset.IVF4096_HNSW128_PQ128-IP-train.faiss \
    --context_encoder_name facebook/dpr-ctx_encoder-multiset-base \
    --index_gpus 2 \
    --gpu_order [2,3,0,1] \
    --indexing_freq 500

# Stop the Ray cluster.
ray stop
```

Since this is an official HF dataset, my reading of the below usage of the `finetune_rag.py` is that I should use the argument `--index_name hf`.

Usage of `finetune_rag.py`:
```
  --index_name INDEX_NAME
                        Name of the index to use: 'hf' for a canonical dataset from the datasets library (default), 'custom' for a local index, or 'legacy' for the orignal one)
  --passages_path PASSAGES_PATH
                        Path to the dataset of passages for custom index. More info about custom indexes in the RagRetriever documentation as well as in `examples/rag/use_own_knowledge_dataset.py`
  --index_path INDEX_PATH
                        Path to the faiss index for custom index. More info about custom indexes in the RagRetriever documentation as well as in `examples/rag/use_own_knowledge_dataset.py`
```
By the wording of the usage, I would think I need only specify `--index_name hf` without `--index_path`. But this results in 
```
Traceback (most recent call last):
  File "/.../rag_end2end/finetune_rag.py", line 820, in <module>
    main(args)
  File "/.../rag_end2end/finetune_rag.py", line 758, in main
    model: GenerativeQAModule = GenerativeQAModule(args)
  File "/.../rag_end2end/finetune_rag.py", line 122, in __init__
    retriever = RagRayDistributedRetriever.from_pretrained(
  File "/.../rag_end2end/distributed_ray_retriever.py", line 159, in from_pretrained
    index = cls._build_index(config)
  File "/.../miniconda3/envs/qa3/lib/python3.9/site-packages/transformers/models/rag/retrieval_rag.py", line 403, in _build_index
    return CanonicalHFIndex(
  File "/fs/nexus-scratch/yzhang42/miniconda3/envs/qa3/lib/python3.9/site-packages/transformers/models/rag/retrieval_rag.py", line 260, in __init__
    raise ValueError("Please provide `index_name` or `index_path`.")
ValueError: Please provide `index_name` or `index_path`.
```

Specifying `--index_name hf` and `--index_path /path/to/.hf_cache/datasets/wiki_dpr/psgs_w100.multiset.compressed/0.0.0/74d4bff38a7c18a9498fafef864a8ba7129e27cb8d71b22f5e14d84cb17edd54/psgs_w100.multiset.IVF4096_HNSW128_PQ128-IP-train.faiss` results in the same behavior.

Specifying Specifying `--index_name custom` and `--index_path [path]` results in trying to load a dummy knowledge base from the `test_finetune.sh` script.
```
Loading passages from /fs/nexus-scratch/yzhang42/rag_end2end/test_run/dummy-kb/my_knowledge_dataset
Traceback (most recent call last):
  File "/.../rag_end2end/finetune_rag.py", line 820, in <module>
    main(args)
  File "/.../rag_end2end/finetune_rag.py", line 758, in main
    model: GenerativeQAModule = GenerativeQAModule(args)
  File "/.../rag_end2end/finetune_rag.py", line 122, in __init__
    retriever = RagRayDistributedRetriever.from_pretrained(
  File "/.../rag_end2end/distributed_ray_retriever.py", line 159, in from_pretrained
    index = cls._build_index(config)
  File "/.../miniconda3/envs/qa3/lib/python3.9/site-packages/transformers/models/rag/retrieval_rag.py", line 397, in _build_index
    return CustomHFIndex.load_from_disk(
  File "/.../miniconda3/envs/qa3/lib/python3.9/site-packages/transformers/models/rag/retrieval_rag.py", line 320, in load_from_disk
    dataset = load_from_disk(dataset_path)
  File "/.../miniconda3/envs/qa3/lib/python3.9/site-packages/datasets/load.py", line 1886, in load_from_disk
    raise FileNotFoundError(f"Directory {dataset_path} not found")
FileNotFoundError: Directory /.../rag_end2end/test_run/dummy-kb/my_knowledge_dataset not found
```
I'm confused as to how this hardcoded value for `dataset_path` even persisted from my run of `test_finetune.sh`


Running `--index_name legacy` results in the follow regardless of specifying `--index_path`
```
OSError: Can't load '/.../.hf_cache/datasets/wiki_dpr/psgs_w100.multiset.compressed/0.0.0/psgs_w100.tsv.pkl'. Make sure that:

- '/.../.hf_cache/datasets/wiki_dpr/psgs_w100.multiset.compressed/0.0.0' is a correct remote path to a directory containing a file named psgs_w100.tsv.pkl

- or '/.../.hf_cache/datasets/wiki_dpr/psgs_w100.multiset.compressed/0.0.0' is the correct path to a directory containing a file named psgs_w100.tsv.pkl.
```

Doing `find /.../.hf_cache/datasets/wiki_dpr -name "psgs_w100.tsv.pkl"` finds nothing.

### Expected behavior

I'm confused as to what the proper arguments should be. Design wise, I think `--index_path=/.../wiki_dpr/psgs_w100.multiset.compressed` should work for my purposes.

I've thought of saving the dataset in question and treating it like the dummy dataset used in `test_finetuning.sh` but I do not have the storage to copy the dataset. I also do not understand how the `dummy-kb` is persisting as an argument in the `finetune_rag_ray_end2end.sh` script.

Comments:
Testing the below hard-coding to get expected behavior:

`finetune_rag.sh` arguments:
```
--index_name custom \
 --index_path /.../.hf_cache/datasets/wiki_dpr/psgs_w100.multiset.compressed/0.0.0/74d4bff38a7c18a9498fafef864a8ba7129e27cb8d71b22f5e14d84cb17edd54/psgs_w100.multiset.IVF4096_HNSW128_PQ128-IP-train.faiss
```

`retrieval_rag.py` hard-coding:
```py
class CustomHFIndex(HFIndexBase):
    def __init__(self, vector_size: int, dataset, index_path=None):
        super().__init__(vector_size, dataset, index_initialized=index_path is None)
        self.index_path = index_path

    @classmethod
    def load_from_disk(cls, vector_size, dataset_path, index_path):
        # print("HERE", dataset_path)
        # print("here2", index_path)
        # raise
        # exit()
        logger.info(f"Loading passages from {dataset_path}")
        if dataset_path is None or index_path is None:
            raise ValueError(
                "Please provide ``dataset_path`` and ``index_path`` after calling ``dataset.save_to_disk(dataset_path)`` "
                "and ``dataset.get_index('embeddings').save(index_path)``."
            )
        # dataset = load_from_disk(dataset_path)
        dataset = load_dataset(path='wiki_dpr', name='psgs_w100.multiset.compressed', split='train') # CHANGED
        return cls(vector_size=vector_size, dataset=dataset, index_path=index_path)
```



This is not a script we maintain.

Yeah, you need to use the custom index.  But you can try to get this to work with the HF index.