Add support for rankllama #294
Conversation
2b78a90
to
dd94777
Compare
| @@ -0,0 +1,91 @@ | |||
| # PECOS XMR Reranker | |||
There was a problem hiding this comment.
Some general comments about the README.md:
- We should also introduce the data schema for training/inference?
- For the Command Line Usage (CLI), we have
pecos.xmr.reranker.train. Should we also havepecos.xmc.reranker.predict? - Do we want to support Python API usage?
There was a problem hiding this comment.
Another high level comment. Can we avoid hard-coded the input columns and make them configurable in the config JSON file? Some hard-coded columns, for example:
- Line 79 of
data_utils.py:keywords - Line 110 of
data_utils.py:contents,titles - Line 296-298 of
model.py:inp_id,ret_idxs,rel
There was a problem hiding this comment.
Another high level comment. Can we avoid hard-coded the input columns and make them configurable in the config JSON file? Some hard-coded columns, for example:
* Line 79 of `data_utils.py`: `keywords` * Line 110 of `data_utils.py`: `contents`, `titles` * Line 296-298 of `model.py`: `inp_id`, `ret_idxs`, `rel`
This can now be specified in the configuration. I added details in the README.md.
Some general comments about the README.md:
* We should also introduce the data schema for training/inference? * For the Command Line Usage (CLI), we have `pecos.xmr.reranker.train`. Should we also have `pecos.xmc.reranker.predict`? * Do we want to support Python API usage?
Added predictions.
| params: The model parameters (RankingModelParams) | ||
| """ | ||
| training_args = train_params.training_args | ||
| training_args.remove_unused_columns = False |
There was a problem hiding this comment.
Is this line still necessary?
There was a problem hiding this comment.
Since we are adding additional information to the output of the collate function, this is needed to avoid it being removed by the trainer.
pecos/xmr/reranker/model.py
Outdated
| """ | ||
| Enable gradient checkpointing for the model | ||
| """ | ||
| self.hf_model.enable_input_require_grads() |
There was a problem hiding this comment.
Not sure if this is the right place to call hf_model.enable_input_require_grads().
From Tevatron RankLlaMA implementation (https://github.com/texttron/tevatron/blob/main/src/tevatron/reranker/modeling.py#L79), they are calling only when both "LoRA" and "training_args.gradient_checkpointing" is enable.
There was a problem hiding this comment.
The following sequence of operation is expected:
- First get_modifed_model is called.
- the trainer calls
gradient_checkpointing_enable, when gradient checkpointing is enabled.
87cc9f6
to
63a083b
Compare
63a083b
to
a648822
Compare
a648822
to
6bfec94
Compare
| lbl_idxs: List[int], | ||
| ): | ||
| """ | ||
| Collate function for training. Tokenizes the input and return features and returns the collated batch. |
There was a problem hiding this comment.
Doc String seems to be out-dated.
pecos/xmr/reranker/model.py
Outdated
| model_params (RankingModel.ModelParams): the model parameters | ||
| train_params (RankingModel.TrainParams): the training parameters | ||
| Returns: | ||
| An instance of UberGlobalModel |
There was a problem hiding this comment.
Doc String seems to be out-dated. Remove UberGlobalModel.
There was a problem hiding this comment.
Updated docstrings
6bfec94
to
725f043
Compare
setup.py
Outdated
| 'transformers>=4.4.2; python_version>="3.9"' | ||
| 'transformers>=4.4.2; python_version>="3.9"', | ||
| 'tqdm>=4.66.4', | ||
| 'peft>=0.11.0', |
There was a problem hiding this comment.
The peft package requires python_version >= 3.8 that's why the latest unit test failed.
Please also check the minimal support python version for other libraries you introduce.
be0c80c
to
6bc4b6a
Compare
6bc4b6a
to
7a3cb6e
Compare
This pull request adds support for RankLlama with LoRA. Detailed instructions are included in
README.md.By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.