NLP Flask API 배포

🎴 들어가며...

이번 포스트에서는 학습된 문장 생성 모델을 Flask API를 통해 배포해보는 과정을 살펴보겠습니다.
해당 과정은 단순 API를 사용하는 방법을 위한 글이므로 아래와 같은 준비 과정이 완료되었다는 가정하에 진행하겠습니다.

• 생성된 문장 생성 모델이 존재
  • .h5 파일로 저장된 weight 사용
• Embedding을 사용하기에 단어 관련 vocab.csv 파일이 존재
  • 아래에서 추가 설명

🥋 과정 설명

해당 배포는 아래의 구조로 이루어진다고 생각하면 됩니다.
1. .h5 의 weight를 받아서 모델을 load
2. vocab.csv 파일을 불러와서 Embedding 된 문장과 연결 준비
3. 모델에 원하는 시작 단어를 입력
4. 생성된 문장을 vocab.csv를 이용해서 한글 문장으로 변환
5. 문장 출력

• vocab.csv는 아래와 같은 형태를 가집니다.
  

🧣 코드 작성

관련 코드는 아래와 같습니다.
참고로, 해당 코드는 두서 없이 작성되어 정리가 되지 않아 보기에 거북(?) 할 수도 있습니다...

from distutils.log import debug
from unicodedata import name
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.layers.experimental.preprocessing import TextVectorization
import os
import re
import string
import random
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers


import io
from torchvision import models
import json
from flask import Flask, jsonify, request
from flask import make_response
import torchvision.transforms as transforms
from PIL import Image

app = Flask(__name__)



temp = pd.read_csv('./park/vocab.csv', encoding='utf-8-sig')
temp.rename(columns={'Unnamed: 0':'vocab'}, inplace=True)
vocab = ['']
for i in range(len(temp)):
  vocab.append(temp.vocab[i])
  
  
embed_dim = 256  # Embedding size for each token
num_heads = 2  # Number of attention heads
feed_forward_dim = 256  # Hidden layer size in feed forward network inside transformer
maxlen = 20
vocab_size = 200000

def create_model():
    inputs = layers.Input(shape=(maxlen,), dtype=tf.int32) 
    # input 정의
    
    embedding_layer = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim) 
    # token Embedding + positional Embedding layer class 정의
    
    x = embedding_layer(inputs) 
    # 선언한 Embedding layer class 이용해 Embedding
    
    transformer_block = TransformerBlock(embed_dim, num_heads, feed_forward_dim) 
    # transformer block layer class 정의
    
    x = transformer_block(x) 
    # 선언한 transformer layer class 이용해 학습
    
    outputs = layers.Dense(vocab_size)(x) 
    # 압축된 결과를 vocab에 맞춰 팽창 (후보단어 선별을 위한 각 단어에 대한 결과치 도출)
    
    model = keras.Model(inputs=inputs, outputs=[outputs, x])  
    # model 정의
    
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) 
    model.compile(
        "adam", loss=[loss_fn, None],
    )  # No loss and optimization based on word embeddings from transformer block
    model.summary()
    return model




class TokenAndPositionEmbedding(layers.Layer):
    def __init__(self, maxlen, vocab_size, embed_dim):
        super(TokenAndPositionEmbedding, self).__init__()
        self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
        self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)

    def get_config(self):
      config = super().get_config().copy()
      config.update({
          'token_emb': self.token_emb,
          'pos_emb': self.pos_emb
      })
      return config

    def call(self, x):
        maxlen = tf.shape(x)[-1]
        # token 위치에 따른 embedding을 하기위함
        positions = tf.range(start=0, limit=maxlen, delta=1)
        positions = self.pos_emb(positions)
        x = self.token_emb(x)
        return x + positions


class TokenAndPositionEmbedding(layers.Layer):
    def __init__(self, maxlen, vocab_size, embed_dim):
        super(TokenAndPositionEmbedding, self).__init__()
        self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
        self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)

    def get_config(self):
      config = super().get_config().copy()
      config.update({
          'token_emb': self.token_emb,
          'pos_emb': self.pos_emb
      })
      return config

    def call(self, x):
        maxlen = tf.shape(x)[-1]
        # token 위치에 따른 embedding을 하기위함
        positions = tf.range(start=0, limit=maxlen, delta=1)
        positions = self.pos_emb(positions)
        x = self.token_emb(x)
        return x + positions

def causal_attention_mask(batch_size, n_dest, n_src, dtype):
    """
    Mask the upper half of the dot product matrix in self attention.
    This prevents flow of information from future tokens to current token.
    1's in the lower triangle, counting from the lower right corner.
    """
    i = tf.range(n_dest)[:, None]
    j = tf.range(n_src)
    m = i >= j - n_src + n_dest
    mask = tf.cast(m, dtype)
    mask = tf.reshape(mask, [1, n_dest, n_src])
    mult = tf.concat(
        [tf.expand_dims(batch_size, -1), tf.constant([1, 1], dtype=tf.int32)], 0
    )
    return tf.tile(mask, mult)


class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
        super(TransformerBlock, self).__init__()
        self.att = layers.MultiHeadAttention(num_heads, embed_dim)
        self.ffn = keras.Sequential(
            [layers.Dense(ff_dim, activation="relu"), layers.Dense(embed_dim),]
        )
        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)
      
    def get_config(self):
      config = super().get_config().copy()
      config.update({
          'att': self.att,
          'ffn': self.ffn,
          'layernorm1': self.layernorm1,
          'layernorm2': self.layernorm2,
          'dropout1':self.dropout1,
          'dropout2':self.dropout2
      })
      return config

    

    def call(self, inputs):
        input_shape = tf.shape(inputs)
        batch_size = input_shape[0]
        seq_len = input_shape[1]
        causal_mask = causal_attention_mask(batch_size, seq_len, seq_len, tf.bool)
        attention_output = self.att(inputs, inputs, attention_mask=causal_mask)
        attention_output = self.dropout1(attention_output)
        out1 = self.layernorm1(inputs + attention_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output)
        return self.layernorm2(out1 + ffn_output)

# 변수 정의 및 모델 함수 정의

embed_dim = 256  # Embedding size for each token
num_heads = 2  # Number of attention heads
feed_forward_dim = 256  # Hidden layer size in feed forward network inside transformer

def create_model():
    inputs = layers.Input(shape=(maxlen,), dtype=tf.int32) 
    # input 정의
    
    embedding_layer = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim) 
    # token Embedding + positional Embedding layer class 정의
    
    x = embedding_layer(inputs) 
    # 선언한 Embedding layer class 이용해 Embedding
    
    transformer_block = TransformerBlock(embed_dim, num_heads, feed_forward_dim) 
    # transformer block layer class 정의
    
    x = transformer_block(x) 
    # 선언한 transformer layer class 이용해 학습
    
    outputs = layers.Dense(vocab_size)(x) 
    # 압축된 결과를 vocab에 맞춰 팽창 (후보단어 선별을 위한 각 단어에 대한 결과치 도출)
    
    model = keras.Model(inputs=inputs, outputs=[outputs, x])  
    # model 정의
    
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) 
    model.compile(
        "adam", loss=[loss_fn, None],
    )  # No loss and optimization based on word embeddings from transformer block
    model.summary()
    return model

word_to_index = {}
for index, word in enumerate(vocab):
    word_to_index[word] = index
index_to_word = {}
for index, word in enumerate(vocab):
    index_to_word[index] = word


# 우리의 리뷰 모델 만들기
class ReviewModel():
  def __init__(self, model, index_to_word, max_tokens=20, top_k=10):
    self.model = model
    self.index_to_word = index_to_word
    self.max_tokens = max_tokens
    self.k = top_k

  def detokenize(self, number):
    return self.index_to_word[number]

  def sample_from(self, logits):
    logits, indices = tf.math.top_k(logits, k=self.k, sorted=True)
    indices = np.asarray(indices).astype("int32")
    preds = keras.activations.softmax(tf.expand_dims(logits, 0))[0]
    preds = np.asarray(preds).astype("float32")
    return np.random.choice(indices, p=preds)

  def tag_to_token(self, tag):
    tokens = [word_to_index.get(_, 1) for _ in tag.split()]
    return tokens

  def tag_to_sentence(self, tag, n=1):
    token = self.tag_to_token(tag)
    res_text = []
    for i in range(n):
      start_tokens = [_ for _ in token]
      num_tokens_generated = 0
      tokens_generated = []
      while num_tokens_generated <= self.max_tokens:
          pad_len = maxlen - len(start_tokens)
          sample_index = len(start_tokens) - 1
          if pad_len < 0:
              x = start_tokens[:maxlen]
              sample_index = maxlen - 1
          elif pad_len > 0:
              x = start_tokens + [0] * pad_len
          else:
              x = start_tokens
          x = np.array([x])
          y, _ = self.model.predict(x)
          sample_token = self.sample_from(y[0][sample_index])
          tokens_generated.append(sample_token)
          start_tokens.append(sample_token)
          num_tokens_generated = len(tokens_generated)
      txt = " ".join(
          [self.detokenize(_) for _ in token + tokens_generated]
      )
      res_text.append(txt)
    return res_text

# 모델 새로 생성하기
park_model = create_model()

# 해당 모델에 저장한 Weight 입히기
park_model.load_weights('./park/park_model.h5')


review_model = ReviewModel(park_model, vocab)

# GET 방식으로 입력하고, 입력을 바탕으로 문장 생성
@app.route('/show')
def show(test=None):
    inp = request.args.get('test')
    output = review_model.tag_to_sentence(inp)
    return(str(output))

if __name__=="__main__":
    app.run(host="0.0.0.0", debug=True)

해당 코드는 /show?test=시작단어 를 통해 GET 방식으로 API에 요청을 보내면 생성된 문장을 화면에 출력하는 Flask API 입니다.

추가로, 제가 머신러닝 쪽을 잘 알지 못해 자세히는 모르지만, 간단하게 전체 과정에 대해 말씀드리면, 아래와 같은 과정으로 이루어집니다.

1. 원하는 문장데이터들을 Embedding으로 변환 후, 모델 학습
2. 생성된 모델을 load
3. 원하는 시작단어를 토큰화
4. 모델에 입력 후, 문장 생성
5. 생성된 문장을 vocab.csv를 이용해서 변환 후, 출력

🧬 응용

현재는 단일 모델로 작성되었지만, 복수의 모델을 사용하는 것도 가능할 것입니다.

단, 주의할 점은 복수의 모델 별로 weight를 불러오거나 vocab.csv를 각각 불러와야 하기에 이를 Flask Server를 켤 때, 한 번에 load 하는 것이 부하를 최소화하는 방법일 것입니다.
한 번에 load를 하는 경우는 일일히 코드를 작성해도 되고, 딕셔너리를 활용해서 코드를 작성하는 것도 가능할 것입니다.

또한, 저는 보여주기 위해 한 개의 python 파일로 모든 함수나 클래스를 작성했지만, 기능에 따라 클래스나 함수를 나누어 작성하고 이를 import하는 형식으로 작성하면 보다 깔끔하게 작성하실 수 있을 것입니다.