FastAPI에서 Huggingface Transformer API를 사용해보자

서론

맞춤법검사기를 제작하면서 Huggingface Transformer API(이하 Transfomer API)를 FastAPI 백엔드에서 사용할 일이 생겼다.

처음에는 당연하게도 이런식으로 구성했다.

@router.post("/correction", response_model=CorrectionResponse)
@limiter.limit("60/seconds")
async def correction(request: Request, correction: CorrectionRequest, user: Session = Depends(get_logged_user)):
    # 1. Tokenize and pad inputs in batches
    batch = [tokenizer([f"{tokenizer.bos_token}{p}{tokenizer.eos_token}"for p in batch], add_special_tokens=True, padding=True, truncation=True, return_tensors="pt",
                       max_length=128) for batch in batches(contents, batch_size)]

    # 2. Initialize empty list for storing generated texts
    generated_texts = []

    # 3. Process batches efficiently using a loop
    for batch_inputs in batch:
        batch_outputs = model.generate(batch_inputs['input_ids'].to(device), max_length=128)
        batch_generated_texts = [tokenizer.decode(output, skip_special_tokens=True) for output in batch_outputs]
        generated_texts.extend(batch_generated_texts)
    return CorrectionResponse(content=generated_texts)

huggingface logo

본론

POST /correction 을 호출할 때마다 다른 백엔드 URL에는 제대로 접속이 되지 않는 현상이 발생한다.
이유는 FastAPI에 있다.
Async IO를 이용하는 FastAPI에서는 이런식으로 오래 걸리는 작업을 백엔드 함수에서 바로 실행한다면 다른 요청에는 제대로 응답하지 못하는 문제가 생길 수 있다. 아마 ASGI를 이용하는 Django에서도 이러한 문제가 생길 수 있다.

결론

우리가 쓰는 asyncio의 이벤트 루프를 이용하자!

loop = asyncio.get_running_loop()

이 구문을 이용하면 loop 변수에는 지금 루프가 아닌 새로운 이벤트루프를 가져올 수 있다. 자세한 내용은 여기를 참고하길 바란다.

이 루프 안에서 Long job를 실행시키자!

result = await loop.run_in_executor(None, run_model, correction.content)

자세한 사용법은 여기에 있다.

이렇게 만들어진 것이 맞춤법 검사기다.