🤗 Transformers를 활용해 Torchserve 배포하기

{
 
'label' : 4,
'text' : '"dr. goldberg offers everything i look for in a general practitioner.
            he's nice and easy to talk to without being patronizing; he's always on
            time in seeing his patients; he's affiliated with a top-notch hospital (nyu)
            which my parents have explained to me is very important in case something happens
            and you need surgery; and you can get referrals to see specialists without
            having to see him first.  really, what more do you need?
            i'm sitting here trying to think of any complaints i have about him,
            but i'm really drawing a blank."'
}

 
from datasets import load_dataset
 
dataset = load_dataset("yelp_review_full")
 
dataset
 
>>> DatasetDict({
    train: Dataset({
        features: ['label', 'text'],
        num_rows: 650000
    })
    test: Dataset({
        features: ['label', 'text'],
        num_rows: 50000
    })
})
 

 
dataset['train'][:3]
 
 
>>> {
    'label': [4, 1, 3],
    'text': [
        "dr. goldberg offers everything i look for in a general practitioner.  he's nice and easy to talk to without being patronizing; he's always on time in seeing his patients; he's affiliated with a top-notch hospital (nyu) which my parents have explained to me is very important in case something happens and you need surgery; and you can get referrals to see specialists without having to see him first.  really, what more do you need?  i'm sitting here trying to think of any complaints i have about him, but i'm really drawing a blank.",
        "Unfortunately, the frustration of being Dr. Goldberg's patient is a repeat of the experience I've had with so many other doctors in NYC -- good doctor, terrible staff.  It seems that his staff simply never answers the phone.  It usually takes 2 hours of repeated calling to get an answer.  Who has time for that or wants to deal with it?  I have run into this problem with many other doctors and I just don't get it.  You have office workers, you have patients with medical needs, why isn't anyone answering the phone?  It's incomprehensible and not work the aggravation.  It's with regret that I feel that I have to give Dr. Goldberg 2 stars.",
        "Been going to Dr. Goldberg for over 10 years. I think I was one of his 1st patients when he started at MHMG. He's been great over the years and is really all about the big picture. It is because of him, not my now former gyn Dr. Markoff, that I found out I have fibroids. He explores all options with you and is very patient and understanding. He doesn't judge and asks all the right questions. Very thorough and wants to be kept in the loop on every aspect of your medical health and your life."
        ]
    }

# train
train_dataset = train_dataset = dataset["train"].shuffle(seed=42).select(range(1000))
 
# validation
validation_dataset = dataset["test"].shuffle(seed=42).select(range(100))

train_dataset.to_csv('data/train.csv',index=False)
validation_dataset.to_csv('data/validation.csv',index=False)
 
# index=False : csv에 포함된 index를 저장하지 않는 기능

# load_dataset
data = load_dataset("csv",data_files={'train':'data/train.csv','validation':'data/test.csv'})
 
data
 
>>>DatasetDict({
   train: Dataset({
       features: ['label', 'text'],
       num_rows: 1000
  })
  validation: Dataset({
      features: ['label', 'text'],
      num_rows: 50
  })
})

 
from transformers import (
        DistilBertForSequenceClassification,
        DistilBertTokenizer,
        Trainer,
        TrainingArguments,
        TrainerCallback
    )
import pandas as pd
 
# load_dataset
data = load_dataset("csv",data_files={'train':'data/train.csv','validation':'data/test.csv'})
 
train_dataset = data['train']
evaluation_dataset = data['validation']
 
print('Complete Loading')
 
 
# Tokenizing
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
 
def tokenize_function(item):
    return tokenizer(item["text"], padding="max_length", max_length=128, truncation=True)
 
train = train_dataset.map(tokenize_function)
evaluation = evaluation_dataset.map(tokenize_function)
 
print('Complete Tokenizing')
 
 
 
tra_args= TrainingArguments(
num_train_epochs=1,
output_dir="test",
logging_steps=10,
# evaluation_strategy="epoch",
)
 
class myCallback(TrainerCallback):
 
    def on_log(self, args, state, control, logs=None, **kwargs):
        print(f'{state.global_step}회 진행 중 ')
 
# num_label = 5
model = DistilBertForSequenceClassification.from_pretrained(
"distilbert-base-uncased", num_labels=5
)
 
trainer = Trainer(
model=model,
args=tra_args,
train_dataset=train,
eval_dataset=evaluation,
callbacks=[myCallback]
)
 
trainer.train()
 

trainer.save_model('torch_model')
tokenizer.save_pretrained('torch_model')

from abc import ABC
import json
import logging
import os
 
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
from ts.torch_handler.base_handler import BaseHandler
 
logger = logging.getLogger(__name__)
 
 
class TransformersClassifierHandler(BaseHandler, ABC):
    def __init__(self):
        super(TransformersClassifierHandler, self).__init__()
        self.initialized = False
 
    def initialize(self, ctx):
        self.manifest = ctx.manifest
 
        properties = ctx.system_properties
        model_dir = properties.get("model_dir")
        self.device = torch.device(
            "cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() else "cpu"
        )
 
        # Read model serialize/pt file
        self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
 
        self.model.to(self.device)
        self.model.eval()
 
        logger.debug("Transformer model from path {0} loaded successfully".format(model_dir))
 
        # Read the mapping file, index to object name
        mapping_file_path = os.path.join(model_dir, "index_to_name.json")
 
        if os.path.isfile(mapping_file_path):
            with open(mapping_file_path) as f:
                self.mapping = json.load(f)
        else:
            logger.warning(
                "Missing the index_to_name.json file. Inference output will not include class name."
            )
 
        self.initialized = True
 
    def preprocess(self, data):
        """Very basic preprocessing code - only tokenizes.
        Extend with your own preprocessing steps as needed.
        """
        print("------- input data 확인 --------")
        print(data)
        text = data[0].get("data")
        if text is None:
            text = data[0].get("body")
 
        logger.info("Received text: '%s'", text)
 
        inputs = self.tokenizer.encode_plus(text, add_special_tokens=True, return_tensors="pt")
        return inputs
 
    def inference(self, inputs):
        """
        Predict the class of a text using a trained transformer model.
        """
        # NOTE: This makes the assumption that your model expects text to be tokenized
        # with "input_ids" and "token_type_ids" - which is true for some popular transformer models, e.g. bert.
        # If your transformer model expects different tokenization, adapt this code to suit
        # its expected input format.
        inputs = inputs.to(self.device)
 
        prediction = self.model(**inputs)[0].argmax().item()
        logger.info("Model predicted: '%s'", prediction)
 
        if self.mapping:
            prediction = self.mapping[str(prediction)]
        return [prediction]
 
    def postprocess(self, inference_output):
        # TODO: Add any needed post-processing of the model predictions here
        logger.info("Model Name: '%s'", self.model.config._name_or_path)
        logger.info("Model predicted: '%s'", inference_output)
        return inference_output
 
 
_service = TransformersClassifierHandler()
 
 
def handle(data, context):
    try:
        if not _service.initialized:
            _service.initialize(context)
 
        if data is None:
            return None
 
        data = _service.preprocess(data)
        data = _service.inference(data)
        data = _service.postprocess(data)
 
        return data
    except Exception as e:
        raise e
 

torch-model-archiver --model-name bert-model --version 1.0 --serialized-file pytorch_model.bin  --handler "handler.py" --extra-files "config.json,vocab.txt"

# MAR_INF 내부 Json 정보
{
    "createdOn": "17/01/2023 18:36:16",
  "runtime": "python",
  "model": {
      "modelName": "bert-model",
    "serializedFile": "pytorch_model.bin",
    "handler": "handler.py",
    "modelVersion": "1.0"
  },
  "archiverVersion": "0.7.0"
}

torchserve --start --model-store model_store --models bert=bert-model.mar

curl -X POST -H "Content-Type: text/plain"  http://127.0.0.1:8080/predictions/bert -d "Stopped back by Mellow Mushroom with my mate Justin from Brew Bros."
 
결과 : 2

#!/bin/bash
 
torchserve --stop
 
model_name="bert-model"
 
# mar 만들기
torch-model-archiver --model-name ${model_name} --version 1.0 --serialized-file pytorch_model.bin  --handler "handler.py" --extra-files "config.json,vocab.txt" --force
 
echo 'MAR File 생성 완료'
 
mv -f ${model_name}.mar model_store
 
실행하기
torchserve --start --model-store model_store --models bert=bert-model.mar --no-config-snapshots