From 04540495f0e8ed8e21cec19bfd64c4213f73836a Mon Sep 17 00:00:00 2001 From: dstnluong-google <129889805+dstnluong-google@users.noreply.github.com> Date: Fri, 14 Jul 2023 11:01:04 -0700 Subject: [PATCH] Add transformers serve files. (#2075) * Add transformers serve files. * Add transformer files to CODEOWNERS * Add image_format_converter.py to utils * Add $ to dockerfile example commands. --- community-content/CODEOWNERS | 1 + .../transformers/dockerfile/serve.Dockerfile | 79 ++++++ .../model_oss/transformers/handler.py | 233 ++++++++++++++++++ .../model_oss/util/image_format_converter.py | 20 ++ 4 files changed, 333 insertions(+) create mode 100644 community-content/vertex_vision_model_garden/model_oss/transformers/dockerfile/serve.Dockerfile create mode 100644 community-content/vertex_vision_model_garden/model_oss/transformers/handler.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/util/image_format_converter.py diff --git a/community-content/CODEOWNERS b/community-content/CODEOWNERS index 6bcd81e00..3e8b0b18f 100644 --- a/community-content/CODEOWNERS +++ b/community-content/CODEOWNERS @@ -12,3 +12,4 @@ /prediction_featurestore_integration @googleapis/vertex-prediction-team /vertex_vision_model_garden/model_oss/util @weigary /vertex_vision_model_garden/model_oss/diffusers @weigary +/vertex_vision_model_garden/model_oss/transformers @dstnluong-google diff --git a/community-content/vertex_vision_model_garden/model_oss/transformers/dockerfile/serve.Dockerfile b/community-content/vertex_vision_model_garden/model_oss/transformers/dockerfile/serve.Dockerfile new file mode 100644 index 000000000..0adbd7fea --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/transformers/dockerfile/serve.Dockerfile @@ -0,0 +1,79 @@ +# Dockerfile for serving dockers for transformers. +# +# To build: +# docker build -f model_oss/transformers/dockerfile/serve.Dockerfile . -t ${YOUR_IMAGE_TAG} +# +# To push to gcr: +# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG} +# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG} +# Switch to this base image for gpu serve. + +FROM pytorch/torchserve:0.7.0-gpu + +USER root + +ENV infer_port=7080 +ENV mng_port=7081 +ENV model_name="transformers_serving" +ENV PATH="/home/model-server/:${PATH}" + +# Install libraries. +ENV PIP_ROOT_USER_ACTION=ignore +RUN python3 -m pip install --upgrade pip +RUN pip install torch==1.13.1 +RUN pip install torchvision==0.14.1 +RUN pip install transformers==4.27.4 +RUN pip install datasets==2.9.0 +RUN pip install accelerate==0.17.0 +RUN pip install triton==2.0.0.dev20221120 +RUN pip install xformers==0.0.16 +RUN pip install google-cloud-storage==2.7.0 +RUN pip install absl-py==1.4.0 + +# Install libraries for document-question-answering. +RUN apt-get update +RUN apt-get install -y --no-install-recommends tesseract-ocr +RUN pip install tesseract==0.1.3 +RUN pip install pytesseract==0.3.10 + +# Install tools. +RUN apt-get update && apt-get install -y --no-install-recommends \ + curl \ + wget \ + vim + +# Copy license. +RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE + +# Copy model artifacts. +COPY model_oss/transformers/handler.py /home/model-server/handler.py +COPY model_oss/util/ /home/model-server/util/ +ENV PYTHONPATH /home/model-server/ + +# Create torchserve configuration file. +RUN echo \ + "default_response_timeout=1800\n" \ + "service_envelope=json\n" \ + "inference_address=http://0.0.0.0:${infer_port}\n" \ + "management_address=http://0.0.0.0:${mng_port}" >> /home/model-server/config.properties + +# Expose ports. +EXPOSE ${infer_port} +EXPOSE ${mng_port} + +# Archive model artifacts and dependencies. +# Do not set --model-file and --serialized-file because model and checkpoint will be dynamically loaded in handler.py. +RUN torch-model-archiver \ +--model-name=${model_name} \ +--version=1.0 \ +--handler=/home/model-server/handler.py \ +--runtime=python3 \ +--export-path=/home/model-server/model-store \ +--archive-format=default \ +--force + +# Run Torchserve HTTP serve to respond to prediction requests. +CMD ["torchserve", "--start", \ + "--ts-config", "/home/model-server/config.properties", \ + "--models", "${model_name}=${model_name}.mar", \ + "--model-store", "/home/model-server/model-store"] diff --git a/community-content/vertex_vision_model_garden/model_oss/transformers/handler.py b/community-content/vertex_vision_model_garden/model_oss/transformers/handler.py new file mode 100644 index 000000000..728bf5edb --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/transformers/handler.py @@ -0,0 +1,233 @@ +"""Custom handler for huggingface/transformers models.""" + +# pylint: disable=g-multiple-import +# pylint: disable=g-importing-member + +import logging +import os +from typing import Any, List, Optional, Tuple + +from PIL import Image +import torch +from transformers import ( + AutoProcessor, + AutoTokenizer, + Blip2ForConditionalGeneration, + Blip2Processor, + BlipForConditionalGeneration, + BlipForQuestionAnswering, + BlipProcessor, + CLIPModel, +) +from transformers import pipeline +from ts.torch_handler.base_handler import BaseHandler + +from util import constants +from util import fileutils +from util import image_format_converter + +DEFAULT_MODEL_ID = "openai/clip-vit-base-patch32" +SALESFORCE_BLIP = "Salesforce/blip" +SALESFORCE_BLIP2 = "Salesforce/blip2" +FLAN_T5 = "flan-t5" +BART_LARGE_CNN = "facebook/bart-large-cnn" + +ZERO_CLASSIFICATION = "zero-shot-image-classification" +FEATURE_EMBEDDING = "feature-embedding" +ZERO_DETECTION = "zero-shot-object-detection" +IMAGE_CAPTIONING = "image-to-text" +VQA = "visual-question-answering" +DQA = "document-question-answering" +SUMMARIZATION = "summarization" +SUMMARIZATION_TEMPLATE = ( + "Summarize the following news article:\n{input}\nSummary:\n" +) + + +class TransformersHandler(BaseHandler): + """Custom handler for huggingface/transformers models.""" + + def initialize(self, context: Any): + """Custom initialize.""" + + properties = context.system_properties + self.map_location = ( + "cuda" + if torch.cuda.is_available() and properties.get("gpu_id") is not None + else "cpu" + ) + self.device = torch.device( + self.map_location + ":" + str(properties.get("gpu_id")) + if torch.cuda.is_available() and properties.get("gpu_id") is not None + else self.map_location + ) + + self.manifest = context.manifest + # The model id is can be either: + # 1) a huggingface model card id, like "Salesforce/blip", or + # 2) a GCS path to the model files, like "gs://foo/bar". + # If it's a model card id, the model will be loaded from huggingface. + self.model_id = ( + DEFAULT_MODEL_ID + if os.environ.get("MODEL_ID") is None + else os.environ["MODEL_ID"] + ) + # Else it will be downloaded from GCS to local first. + # Since the transformers from_pretrained API can't read from GCS. + if self.model_id.startswith(constants.GCS_URI_PREFIX): + gcs_path = self.model_id[len(constants.GCS_URI_PREFIX) :] + local_model_dir = os.path.join(constants.LOCAL_MODEL_DIR, gcs_path) + logging.info("Download %s to %s", self.model_id, local_model_dir) + fileutils.download_gcs_dir_to_local(self.model_id, local_model_dir) + self.model_id = local_model_dir + + self.task = ( + ZERO_CLASSIFICATION + if os.environ.get("TASK") is None + else os.environ["TASK"] + ) + logging.info( + "Handler initializing task:%s, model:%s", self.task, self.model_id + ) + + if SALESFORCE_BLIP in self.model_id: + # pipeline() hasn't been ready for Salesforce/blip models. + self.salesforce_blip = True + self._create_blip_model() + else: + self.salesforce_blip = False + if self.task == FEATURE_EMBEDDING: + self.model = CLIPModel.from_pretrained(self.model_id).to( + self.map_location + ) + self.tokenizer = AutoTokenizer.from_pretrained(self.model_id) + self.processor = AutoProcessor.from_pretrained(self.model_id) + elif self.task == SUMMARIZATION and FLAN_T5 in self.model_id: + self.pipeline = pipeline( + task=self.task, + model=self.model_id, + torch_dtype=torch.bfloat16, + device_map="auto", + ) + else: + self.pipeline = pipeline( + task=self.task, model=self.model_id, device=self.device + ) + + self.initialized = True + logging.info("Handler initialization done.") + + def _create_blip_model(self): + """A helper for creating BLIP and BLIP2 models.""" + if SALESFORCE_BLIP2 in self.model_id: + self.torch_type = torch.float16 + self.processor = Blip2Processor.from_pretrained(self.model_id) + self.model = Blip2ForConditionalGeneration.from_pretrained( + self.model_id, torch_dtype=self.torch_type + ).to(self.map_location) + else: + self.torch_type = torch.float32 + self.processor = BlipProcessor.from_pretrained(self.model_id) + if self.task == IMAGE_CAPTIONING: + self.model = BlipForConditionalGeneration.from_pretrained( + self.model_id + ).to(self.map_location) + elif self.task == VQA: + self.model = BlipForQuestionAnswering.from_pretrained(self.model_id).to( + self.map_location + ) + + def _reformat_detection_result(self, data: List[Any]) -> List[Any]: + """Reformat zero-shot-object-detection output.""" + if not data: + return [data] + boxes = {} + boxes["label"] = data[0]["label"] + boxes["boxes"] = [] + for item in data: + box = {} + box["score"] = item["score"] + box.update(item["box"]) + boxes["boxes"].append(box) + outputs = [boxes] + return outputs + + def preprocess( + self, data: Any + ) -> Tuple[Optional[List[str]], Optional[List[Image.Image]]]: + """Preprocess input data.""" + texts = None + images = None + if "text" in data[0]: + texts = [item["text"] for item in data] + if "image" in data[0]: + images = [ + image_format_converter.base64_to_image(item["image"]) for item in data + ] + return texts, images + + def inference(self, data: Any, *args, **kwargs) -> List[Any]: + """Run the inference.""" + texts, images = data + preds = None + if self.task == ZERO_CLASSIFICATION: + preds = self.pipeline(images=images, candidate_labels=texts) + elif self.task == ZERO_DETECTION: + # The object detection pipeline doesn't support batch prediction. + preds = self.pipeline(image=images[0], candidate_labels=texts[0]) + elif self.task == IMAGE_CAPTIONING: + if self.salesforce_blip: + inputs = self.processor(images[0], return_tensors="pt").to( + self.map_location, self.torch_type + ) + preds = self.model.generate(**inputs) + preds = [ + self.processor.decode(preds[0], skip_special_tokens=True).strip() + ] + else: + preds = self.pipeline(images=images) + elif self.task == VQA: + # The VQA pipelines doesn't support batch prediction. + if self.salesforce_blip: + inputs = self.processor(images[0], texts[0], return_tensors="pt").to( + self.map_location, self.torch_type + ) + preds = self.model.generate(**inputs) + preds = [ + self.processor.decode(preds[0], skip_special_tokens=True).strip() + ] + else: + preds = self.pipeline(image=images[0], question=texts[0]) + elif self.task == DQA: + # The DQA pipelines doesn't support batch prediction. + preds = self.pipeline(image=images[0], question=texts[0]) + elif self.task == FEATURE_EMBEDDING: + preds = {} + if texts: + inputs = self.tokenizer( + text=texts, padding=True, return_tensors="pt" + ).to(self.map_location) + text_features = self.model.get_text_features(**inputs) + preds["text_features"] = text_features.detach().cpu().numpy().tolist() + if images: + inputs = self.processor(images=images, return_tensors="pt").to( + self.map_location + ) + image_features = self.model.get_image_features(**inputs) + preds["image_features"] = image_features.detach().cpu().numpy().tolist() + preds = [preds] + elif self.task == SUMMARIZATION and FLAN_T5 in self.model_id: + texts = [SUMMARIZATION_TEMPLATE.format(input=text) for text in texts] + preds = self.pipeline(texts, max_length=130) + elif self.task == SUMMARIZATION and self.model_id == BART_LARGE_CNN: + preds = self.pipeline( + texts[0], max_length=130, min_length=30, do_sample=False + ) + else: + raise ValueError(f"Invalid TASK: {self.task}") + return preds + + def postprocess(self, data: Any) -> List[Any]: + if self.task == ZERO_DETECTION: + data = self._reformat_detection_result(data) + return data diff --git a/community-content/vertex_vision_model_garden/model_oss/util/image_format_converter.py b/community-content/vertex_vision_model_garden/model_oss/util/image_format_converter.py new file mode 100644 index 000000000..4f5e59526 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/util/image_format_converter.py @@ -0,0 +1,20 @@ +"""Image format converter util lib.""" + +import base64 +import io + +from PIL import Image + + +def image_to_base64(image: Image.Image) -> str: + """Convert a PIL image to a base64 string.""" + buffer = io.BytesIO() + image.save(buffer, format="JPEG") + image_str = base64.b64encode(buffer.getvalue()).decode("utf-8") + return image_str + + +def base64_to_image(image_str: str) -> Image.Image: + """Convert a base64 string to a PIL image.""" + image = Image.open(io.BytesIO(base64.b64decode(image_str))) + return image