Rust에서 쓸 수 있는 NLP 라이브러리

서론

개인용 노트 관리 툴에서 제공하는 연관문서 검색 기능 때문에 빌드가 자주 깨져서 대안을 찾아보았다.

Rust-bert

Rust-bert는 딱 LLM만 빼고 나머지 짜잘한 모델들을 쉽게 쓸 수 있게 만들어놓은 라이브러리라고 할 수 있다. 여기서 LLM만 되었어도 내가 llm 전문 라이브러리같은걸 만들 필요도 없었을 것이다.

다만 이 라이브러리가 libtorch를 필요로 하는데, libtorch의 버전과 거기에 해당하는 tch-rs의 버전이 딱 맞아야지만 정상적으로 구동되는 단점이 있다. 이게 의외로 큰데, libtorch가 대충 몇달에 한번꼴로 활발하게 버전업이 되는 라이브러리라서 꽤나 자주 망가지는 편이고, 내가 쓰는 툴이 맥과 리눅스에서 컴파일이 되어야 하는데 두 OS가 지원하는 버전이 달라질 때가 잦고 그때마다 삽질을 많이 해야했다.

어쩔 수 없이 대안을 찾다가, candle을 발견하게 되었다.

Candle

Candle은 libtorch같은 C 라이브러리 의존 없이 Rust만으로 구현한… 뭐랄까 libtorch 같은 느낌의 라이브러리? 라고 할 수 있다. 다만 모델을 다운로드받고 구동하는 수준의 구현에 집중하고 있어서 Python에서 쓸 수 있는 라이브러리와 비교하면 기능이 크게 부족하다 할 수 있고, 심지어 rust-bert와 비교해도 많이 부족하다. 그러나 내가 원하는 기능이 단순 문서 embedding 계산 및 KeyBERT로 비교적 단순한 편이고, Pure rust로만 구현되어 있어서 개발환경 관리가 편하기 때문에 한번 찍먹해보려고 했다.

단순 모델 구동하기

원래 내가 쓰던 문서 요약 모델은 all-MiniLM-L12-v2 였지만 all-MiniLM-L6-v2로 변경했다. 왜냐하면 Candle로 구현한 구동 예제 가 있어서 재연이 쉬울 것 같았기 때문.

일단 rust-bert와 candle의 계산 결과가 같게 나오는 것을 확인하고 다음 단계로 이동.

KeyBERT 구현하기

대충 KeyBERT의 구현은 다음과 같다.

  1. 문서의 임베딩을 계산
  2. 단어 후보를 추리기 위해 문서를 단어 단위로 쪼개고 distintct set을 계산. 원래는 ngram을 지원하지만 내가 처음부터 1단어만 써봤기 때문에 간단하게 1단어만 쓰기로 했다.
  3. 모든 단어 후보의 임베딩을 계산
  4. 문서와 단어 후보의 임베딩 간 거리를 계산
  5. 가장 유사도가 높은 녀석부터 순서대로 5마리 컷

Rust-bert의 구현을 참조해보니 n개의 단어 모두를 tokenizer에 넣고 돌려서 하나의 tensor에서 계산함으로써 병렬처리를 구현했다. (즉, n개의 단어 모두의 token list를 계산해서 하나의 n x m tensor(여기서 m은 각 단어의 token list중 가장 긴 것의 길이)를 만든다는 뜻)

이렇게 하면 임베딩 계산도 한번의 모델 forward로 계산 가능하고, 거리 계산도 cosine similarity 기준 한번의 matmul 연산으로 계산가능하기 때문에 빠르게 할 수 있구나… 아 배웠다… 하면서 계산해봤는데,

결과가 달라?!

망해버림

디버깅을 해 보니까 일단 문서의 임베딩 계산부터 뭔가 달라졌는데, 토큰 목록 뒤에 0을 여러개 붙이는 경우 임베딩 계산결과가 달라지는 것을 발견했다. 원래는 마스킹을 해서 안쓰는 토큰이라고 표시를 해 줘야 하는데 그걸 지원하지 않는 것. 최근에 기능을 추가했는데 아직 릴리즈 버전에 추가는 안됐다. (그리고 아직 100% 완벽하게 돌아가진 않는 것 같기도 하다)

따라서 임베딩 계산을 할때 token 목록 갯수를 정확하게 잘라줘야 하는데, 이러면 병렬처리가 안돼잖아!

망해버림2

병렬처리가 안되니 울며 겨자먹기로 단어 하나마다 임베딩 계산을 하고 그것들을 합쳐서 단어 목록 임베딩을 만들어내는 방식으로 땜질처리를 했다. 상대적으로 매우 느려지지만 일단 기능이 되도록 이전하는 게 목표라서 별 수 없었다.

망해버림3

어찌어찌 기능 다 되게 이전했는데, 막상 기존 문서에서 키워드 추출 기능을 돌려보니 작은 문서에선 의도대로 돌아가는데 조금 문서길이가 커지니까 여지없이 panic이 발생하는 것을 확인했다. 여기쯤에서 더 쓸 수 있는 시간이 없어서 타임아웃 당하고 다음에 다시 디버깅해보기로 했다.

그나마 잘된 것

내 코드의 대부분 구현을 가능한 한 trait을 이용해서 구현했었는데, 이번에 이전할 때 큰 도움이 되었다. 기존 코드와 동일한 동작을 하도록 trait으로 맞춰주니까 다른 부분을 추가로 건드릴 필요 없이 동작이 돼서 기부니가 좋았다.

결론

성능때문에 rust 쓴다고 했는데 막상 빌드가 불안하니까 성능이고 뭐고 집어치우고 어떻게든 컴파일만 되게 만든 개똥코드를 돌리게 되었다. 이래도 Rust 쓰쉴?