OpenNMT / CTranslate2

Fast inference engine for Transformer models

Home Page:https://opennmt.net/CTranslate2

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

[Feature Request] add a classification / estimation layer on top of Encoders

vince62s opened this issue · comments

This has been a long time request.
Recently resurfaced with #1534 .
#1342 is also an example on how to perform classification but HF library is needed.

Technically this is not very difficult since it is just a matter of adding some extra Linear layer after the Encoder.
However this can take various form since all features are not identical.

I would like two use-cases to work to make sure we can check the output.

Use-case 1: XLM-Roberta classifier
it is just a matter of adding a new "Layer" (in CT2 taxonomy) named "classifier" which replicates this: https://github.com/huggingface/transformers/blob/main/src/transformers/models/xlm_roberta/modeling_xlm_roberta.py#L1167
which you can see here: https://github.com/huggingface/transformers/blob/main/src/transformers/models/xlm_roberta/modeling_xlm_roberta.py#L1434-L1453
is a very simple module.

Use-case 2: Unbabel Comet Estimator
This is again a very simple add-on.
A new "Layer" named "estimator" needs to be added based on this:
https://github.com/Unbabel/COMET/blob/master/comet/modules/feedforward.py#L26-L72
since we just need to use it for inference there are only 3 nn.Linear named:
estimator.ff.0
estimator.ff.3
estimator.ff.6
because 1, 2, 4, 5 are not learned weights.

Of course the converter for https://github.com/OpenNMT/CTranslate2/blob/master/python/ctranslate2/converters/transformers.py#L1970
would need to be modified to reflect the new layer.

If everything works fine for XLM-Roberta we can eventually adapt for other encoder-based models.

LMK if you need more info to implement.

That would be great to be implemented!