[ML] HuggingFace 맛보기(Translator)

개요

이번에 HuggingFace 플랫폼에 대해 알게되었다. HuggingFace의 쉬운 interface을 직접 사용해보기로 했다.

모델

사용한 모델은 facebook/nllb-200-distilled-600M이다. 다음과 같은 코드로 간단하게 번역기를 사용해볼 수 있다.

# import hugginface model
from transformers import pipeline, AutoTokenizer
import torch
import time

model_checkpoint = "facebook/nllb-200-distilled-600M"
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# define tokenizer that allows long token length
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, max_length=2048)


translator = pipeline(task="translation_en_to_kr", model=model_checkpoint, device=device, tokenizer=tokenizer)

# use translator
# text = "Lockheed Martin Delivers Initial 5G Testbed To U.S. Marine Corps And Begins Mobile Network Experimentation"
text = "“Dad, these glasses can help correct your red-green color blindness,” said Hailey. He slowly put them on, and stared at the birthday presents on the table. Seeing vivid red and green colors for the first time ever, he started to cry. “Incredible! Look at those wonderful colors!” He shouted in amazement."
start = time.time()
output = translator(text, src_lang="eng_Latn", tgt_lang="kor_Hang", max_length=2048)
end = time.time()
print(output)

print(f'time taken: {end - start}s')

Result:
[{'translation_text': ' 아빠, 이 안경은 당신의 붉은 녹색 색의 맹목을 바로잡는데 도움이 될 수 있다.  해일리는 말했다. 그는 천천히 그것을 착용하고 테이블에 있는 생일 선물들을 쳐다봤다. 생생한 빨간색과 녹색 색을 처음으로 보고, 그는 울기 시작했다.  믿을 수 없을 정도로! 그 멋진 색을 보세요! 그는 놀라서 외쳤다.'}]
time taken: 1.921203374862671s

CPU를 사용했을 때 2초가 소요되었고, GPU를 사용하면 0.3초가 소요되었다.
모델의 크기는 대략 2.4GB였기 때문에, 웬만한 GPU로도 inference을 할 수 있을 것이다.

Finetuning

사실 위에서 번역한 문장도 어색한 것을 확인할 수 있다. HuggingFace는 기존의 pretrained 모델을 finetuning할 수 있게 해준다. 최근에는 법률 AI같이 특정 분야에 대한 번역 능력을 확 끌어올리도록 finetuning된 AI가 등장하고 있다.

나중에 기회가 된다면 흥미로운 분야에 대해 LLM을 finetuning해볼 기회가 생기면 좋겠다.