ChatGPT API 토큰 수대로 자르기 (javascript)

tilnote2023-09-14 11:58

이번에는 텍스트를 받아서 토큰 수대로 자르는 기능을 만들어 보겠습니다. 가장 기본적인 기능이고 많이 사용하는 기능이 아닐까 생각해 봅니다. 예) 3000 토큰으로 자르기.

환경은 node.js 서버입니다.

먼저 OpenAI에서 만든 tiktoke을 js 버전으로 포팅한 @dqbd/tiktoken을 사용합니다. @dqbd/tiktoken

npm install @dqbd/tiktoken

또는

yarn add @dqbd/tiktoken

모든 encoder가 들어있는 것을 import 할 수 도 있지만 GPT-3.5와 GPT-4에서 사용하는 cl100k base만 import 해서 사용하겠습니다.

const { Tiktoken } = require("@dqbd/tiktoken/lite");
const cl100k_base = require("@dqbd/tiktoken/encoders/cl100k_base.json");

const encoding = new Tiktoken(
  cl100k_base.bpe_ranks,
  cl100k_base.special_tokens,
  cl100k_base.pat_str
);
const tokens = encoding.encode("hello world");
encoding.free();

이렇게 한 후 출력을 해보면 글자가 토큰화가 되어 숫자의 배열이 출력됩니다. 그리고 encoding을 한 후에 메모리가 많이 소모됨으로 encoding.free()를 꼭꼭 해주세요.

Uint32Array(2) [ 15339, 1917 ]

2개의 토큰을 가진 배열임을 알 수 있죠? 이제 이걸 토큰 수 대로 잘라서 사용하면 됩니다. 예를 들어서 특정 텍스트를 10개의 토큰 단위로 자르는 함수를 만들어 봅시다.

  const tokens = encoding.encode(text);
  let chunked = [];
  const n = 10;
  const overwrap = 0;
  for (let i = 0; i < tokens.length; i += n - overwrap) {
    const slice = tokens.slice(i, i + n);
    chunked.push(slice);
  }
  encoding.free();
  return chunked;

tokens을 array로 받아서 n개의 토큰으로 잘라주는 함수입니다. for 구문에서 i를 n개씩 증가시켜주면됩니다. 만약 겹치는 부분을 허용하려면 overwrap을 지정해 주면 됩니다.

[
  Uint32Array(10) [
    61415, 16969, 10997,
      228, 58260,   223,
      108, 18359, 65677,
    16306
  ],
  Uint32Array(7) [
      112, 16969,
    90960, 71682,
    25941, 80052,
       13
  ]
] chunked

이런식으로 잘 잘린 것을 알 수 있죠.

하지만 우리가 원하는 것은 토큰화된 텍스트가 아니라 텍스트입니다. 텍스트로 decode한 후 반환하는 기능을 추가해봅시다.

  for (let i = 0; i < tokens.length; i += n - overwrap) {
    const slice = tokens.slice(i, i + n);
    let decoded = new TextDecoder().decode(encoding.decode(slice));
    chunked.push(decoded);
  }

[ '나는 토큰을 자�', '�는 서비스입니다.' ]

이런식으로 잘라서 반환하는 것을 볼 수 있습니다. 한글의 경우 2~3토큰을 차지하기 때문에 토큰 수대로 자르면 위처럼 ? 로 처리할 수 없는 유니코드 처리가 되는 경우가 많습니다. 어쩔 수 없다고 생각합니다. 어차피 잘린 상태의 텍스트도 인식할 수 있게 만들어진 LLM이기에 이번에는 물음표 문자를 없애주는 처리만 하도록 하겠습니다.

decoded = decoded.replace(/�/g, "");

전체 함수는 이렇게 생겼습니다.

function cutToken(text) {
  const encoding = new Tiktoken(
    cl100k_base.bpe_ranks,
    cl100k_base.special_tokens,
    cl100k_base.pat_str
  );
  const tokens = encoding.encode(text);
  let chunked = [];
  const n = 10;
  const overwrap = 0;
  for (let i = 0; i < tokens.length; i += n - overwrap) {
    const slice = tokens.slice(i, i + n);
    let decoded = new TextDecoder().decode(encoding.decode(slice));
    decoded = decoded.replace(/�/g, "");
    chunked.push(decoded);
  }
  encoding.free();
  return chunked;
}

[ '나는 토큰을 자', '는 서비스입니다.' ]

자 이제 자바스크립트에서 토큰 수대로 마음대로 잘라서 사용할 수 있겠죠? 글자 수로 잘라도 되지만 어디까지나 임시 방편으로 토큰 수로 정확하게 자르는 것이 필요하게 될 것입니다. LLM에 있어 토큰 자르기는 필수적인 기능이라고 생각합니다.

여러분이 더욱 발전시켜 보세요~!