streamlit을 통해 NLP 모델을 배포해 보자

모델을 다 훈련시키고 테스트까지 완료 했더니 저장된 모델의 크기가 1.3기가가 넘었다.

당연히 깃헙에 그대로 업로드 할 수 없었고 그대로 다른 방법을 알아봐야 하나 싶었지만,

묘수가 있었기에 여기에 정리한다.

1. S3에 훈련된 모델을 업로드 한다.
2. @st.cache_resource을 이용해 다운로드 후 변수에 저장하는 함수를 만든다.

예를 들자면, huggingface에서 모델과 토크나이저를 다운 받고 로컬에 저장해 두고 싶다면 아래와 같은 코드를 넣으면 된다.

@st.cache_resource
def load_roberta_model(token):
    tokenizer = RobertaTokenizer.from_pretrained('roberta-base', token=token)
    roberta_model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=4, token=token)
    return tokenizer, roberta_model

그리고 사전 훈련된 모델을 불러오고 싶다면 아래와 같이 코드를 작성하면 된다.

@st.cache_resource
def load_model(device):
    s3 = boto3.client(
        's3',
        aws_access_key_id=boto_access_key,
        aws_secret_access_key=boto_secret_key
    )
    s3.download_file(BUCKET_NAME, BUCKET_KEY, TARGET_NAME)
    model = RoBERTaFineTuner.load_from_checkpoint(TARGET_NAME
, model=roberta_model, device=device.type)
    return model

그리고 아무 문제 없었다는 듯이 아래와 같이 함수를 로딩해 주면 된다.

tokenizer, roberta_model = load_roberta_model(token)

device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')

model = load_model(device)

predict 함수의 경우

def predict(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    with torch.no_grad():
        logits = model(**inputs).logits
    predicted_class_id = logits.argmax().item()
    return label_map[predicted_class_id]

다음과 같이 작성하면 streamlit에서도 아무 문제 없이 실행 할 수 있다.