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.
2026-09-26 14:42:04 +00:00 · 2023-07-14 18:01:04 +00:00
parent ad6e554984
commit 04540495f0
4 changed files with 333 additions and 0 deletions
@@ -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
@@ -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"]
@@ -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
@@ -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