search

neulab / knn-transformers

PyTorch + HuggingFace code for RetoMaton: "Neuro-Symbolic Language Modeling with Automaton-augmented Retrieval" (ICML 2022), including an implementation of kNN-LM and kNN-MT
MIT License
271 stars 22 forks source link

How to apply knn-transformers to a custom pretrained machine translation model? #11

Closed vhientran closed 1 year ago

vhientran commented 1 year ago

Hi @urialon ,

Thank you very much for releasing the source code which applies kNN for the Machine Translation task. However, only pre-trained models available on the Huggingface hub seem valid. Recently, I designed and trained a custom NMT model for the Low-Resource Machine Translation task using PyTorch and Transformers libraries. In the inference phase, I used the method: my_pretrained_model.generate(input_ids) as the traditional translation to translate a testing input sentence. Using your source code, I want to apply kNN for my_pretrained_model in the inference phase. However, since my model is custom, I wonder how to do that. Please guide me in applying your source code to use kNN for my custom pre-trained NMT model. Many thanks for your help!

urialon commented 1 year ago

Hi @vhientran , Thank you for your interest in our work!

I think that there is no difference, and you can use your own model as well. The only place where you will need to make a change is this dictionary: https://github.com/neulab/knn-transformers/blob/master/knnlm.py#L254C31-L271

that tells the model where to find the last decoder layer, according to the chosen architecture.

Let me know if you encounter any problems.

Best, Uri

vhientran commented 1 year ago

Thank you very much for your quick reply. I will try it and let you know soon. Thanks!