nllb200을 이용한 다국어 번역

페이스북에서는 다국어 번역기 모델 nllb를 공개한 적이 있는데 다양한 파라미터와 증류된 모델까지 공개해놨다. 

https://huggingface.co/facebook/nllb-200-distilled-600M

facebook/nllb-200-distilled-600M · Hugging Face

 

로컬PC에서는 distilled 600M이 GPU 사용량도 적어서 이걸 사용하면된다.

다국어 번역이라서 한국어 번역에는 어려운 걸 번역시키면 다른 언어 토큰이 나오는 좀 약간 애매한 점이 있었는데, 원본 증류되지 않은 모델은 써보질 못해서 어떤지 모른다.

 

사용하기는 아주 쉽게 pipline을 이용해서 사용하면된다.

여기서 중요한 것은 src_lang과 tgt_lang인데 src가 번역할 언어고 tgt가 번역될 언어이다.

영어 → 한글을 할려면,  eng_Latn을 src에, kor_Hang을 tgt에 넣어준다.

 

이런 언어코드를 뭘 써야될지 몰라서 좀 고통 받았는데 special_tokens_map.json에 들어 있긴한데 이름을보고 추측해서 써야된다.

 

주로 언어의 3글자와 문자4글자로 구성된 단어인거 같다.

귀찬게 국가만 넣으면 미리 바꿔주는 코드도 만들어봤었다. 

 

하나의 모델에 200개 언어가 가능하다니 나름 유용하게 사용가능하다.

import os
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

model_path: str = os.path.join(root_path, f"llm_model/nllb-200-distilled-600M/")

src_lang = convert_language_code(src_lang)
tgt_lang = convert_language_code(tgt_lang)

model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer, max_length=512,
                                src_lang=src_lang, tgt_lang=tgt_lang, device=0,
                                repetition_penalty=1.2, no_repeat_ngram_size=3, num_beams=4, num_beam_groups=4,
                                penalty_alpha=0.6, top_k=4)