Use Gemma to gauge sentiment and summarize conversations

ARG SERVING_BUILD_IMAGE=tensorflow/tensorflow:2.16.1-gpu

FROM ${SERVING_BUILD_IMAGE}
WORKDIR /workspace

COPY gemma2  gemma2
RUN apt-get update -y && apt-get install -y cmake && apt-get install -y vim

apache_beam[gcp]==2.54.0
keras_nlp==0.14.3
keras==3.4.1
jax[cuda12]

FROM <DOCKERFILE_IMAGE>

WORKDIR /workspace

# Copy files from the official SDK image, including the script and dependencies.
COPY --from=apache/beam_python3.11_sdk:2.54.0 /opt/apache/beam /opt/apache/beam


COPY requirements.txt requirements.txt
RUN pip install --upgrade --no-cache-dir pip \
    && pip install --no-cache-dir -r  requirements.txt

# Copy the model directory downloaded from Kaggle and the pipeline code.
COPY pipeline.py pipeline.py

# The colab was tested and run with a JAX backend to let Dataflow workers
# pick up the environment needed to include in the Env of the image.
ENV KERAS_BACKEND="jax"
ENV XLA_PYTHON_CLIENT_MEM_FRACTION="0.9"


# Set the entrypoint to the Apache Beam SDK launcher.
ENTRYPOINT ["/opt/apache/beam/boot"]

%pip install apache_beam[gcp]=="2.54.0" keras_nlp=="0.14.3" keras=="3.5.0" jax[cuda12]

import os

os.environ["KERAS_BACKEND"] = "jax"  # Or "tensorflow" or "torch".
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.9"

import keras
import keras_nlp
import numpy as np
import json
import ast
import re
import logging

import apache_beam as beam
from apache_beam.ml.inference import utils
from apache_beam.ml.inference.base import RunInference
from apache_beam.options import pipeline_options
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import SetupOptions
from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.options.pipeline_options import WorkerOptions
from apache_beam.ml.inference import utils
from apache_beam.ml.inference.base import ModelHandler
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.ml.inference.base import KeyedModelHandler
from keras_nlp.models import GemmaCausalLM
from typing import Any, Dict, Iterable, Optional, Sequence

options = PipelineOptions()
options.view_as(StandardOptions).streaming = True
options.view_as(SetupOptions).save_main_session = True

# options.view_as(StandardOptions).runner = "dataflowrunner"
# options.view_as(GoogleCloudOptions).project = <PROJECT>
# options.view_as(GoogleCloudOptions).temp_location= <TMP LOCATION>
# options.view_as(GoogleCloudOptions).region= "us-west1"
# options.view_as(WorkerOptions).machine_type= "g2-standard-4"
# options.view_as(WorkerOptions).worker_harness_container_image = <IMAGE YOU BUILT>
# options.view_as(WorkerOptions).disk_size_gb=200
# options.view_as(GoogleCloudOptions).dataflow_service_options=["worker_accelerator=type:nvidia-l4;count:1;install-nvidia-driver"]

topic_reviews="<PubSub Topic>"

class GemmaModelHandler(ModelHandler[str,
                                     PredictionResult,GemmaCausalLM
                                     ]):
    def __init__(
        self,
        model_name: str = "",
    ):
        """ Implementation of the ModelHandler interface for Gemma using text as input.

        Example Usage::

          pcoll | RunInference(GemmaModelHandler())

        Args:
          model_name: The Gemma model uri.
        """
        self._model_name = model_name
        self._env_vars = {}
    def share_model_across_processes(self)  -> bool:
        """Returns whether to share a single model in memory across processes.

        This is useful when the loaded model is large, preventing potential
        out-of-memory issues when running the pipeline.

        Returns:
          bool
        """
        return True

    def load_model(self) -> GemmaCausalLM:
        """Loads and initializes a model for processing."""
        return keras_nlp.models.GemmaCausalLM.from_preset(self._model_name)

    def run_inference(
        self,
        batch: Sequence[str],
        model: GemmaCausalLM,
        inference_args: Optional[Dict[str, Any]] = None
    ) -> Iterable[PredictionResult]:
        """Runs inferences on a batch of text strings.

        Args:
          batch: A sequence of examples as text strings.
          model:
          inference_args: Any additional arguments for an inference.

        Returns:
          An Iterable of type PredictionResult.
        """
        # Loop each text string, and use a tuple to store the inference results.
        predictions = []
        for one_text in batch:
            result = model.generate(one_text, max_length=1024)
            predictions.append(result)
        return utils._convert_to_result(batch, predictions, self._model_name)

prompt_template = """
<prompt>
Provide the results of doing these two tasks on the chat history provided below for the user {}
task 1 : assess if the tone is happy = 1 , neutral = 0 or unhappy = -1
task 2 : summarize the text with a maximum of 512 characters
Return the answer as a JSON string with fields [sentiment, summary] do NOT explain your answer

@@@{}@@@
<answer>
"""
chat_text = """
id 221: Hay I am really annoyed that your menu includes a pizza with pineapple on it!
id 331: Sorry to hear that , but pineapple is nice on pizza
id 221: What a terriable thing to say! Its never ok, so unhappy right now!
"""

# Example input
chat =json.dumps({"id" : 42, "user_id" : 221 , "chat_message" : chat_text})
print(chat)

keyed_model_handler = KeyedModelHandler(GemmaModelHandler('gemma_2'))

# Create the prompt by using the information from the chat.
class CreatePrompt(beam.DoFn):
  def process(self, element, *args, **kwargs):
    user_chat = json.loads(element)
    chat_id = user_chat['id']
    user_id = user_chat['user_id']
    messages = user_chat['chat_message']
    yield (chat_id, prompt_template.format(user_id, messages))

def extract_model_reply(model_inference):
    print(model_inference)
    match = re.search(r"(\{[\s\S]*?\})", model_inference)
    json_str = match.group(1)
    print(json_str)
    result = json.loads(json_str)
    if all(key in result for key in ['sentiment', 'summary']):
        return result
    raise Exception('Malformed model reply')

    # @title
class SentimentAnalysis(beam.DoFn):
    def process(self, element):
        key = element[0]
        match = re.search(r"@@@([\s\S]*?)@@@", element[1].example)
        chats = match.group(1)

        try:
            # The result contains the prompt. Replace the prompt with "".
            result = extract_model_reply(element[1].inference.replace(element[1].example, ""))
            processed_result = (key, chats, result['sentiment'], result['summary'])

            if (result['sentiment'] ==-1):
              output = beam.TaggedOutput('negative', processed_result)
            else:
              output = beam.TaggedOutput('main', processed_result)

        except Exception as err:
            print("ERROR!" + str(err))
            output = beam.TaggedOutput('error', element)

        yield output

gemma_inference = RunInference(keyed_model_handler)

with beam.Pipeline(options=options) as p:
  chats = (p | "Read Topic" >>
            beam.io.ReadFromPubSub(subscription=topic_reviews)
            | "Parse" >> beam.Map(lambda x: x.decode("utf-8")))
  prompts = chats | "Create Prompt" >> beam.ParDo(CreatePrompt())
  results = prompts | "RunInference-Gemma" >> gemma_inference
  filtered_results = results | "Process Results" >> beam.ParDo(SentimentAnalysis()).with_outputs('main','negative','error')
  generated_responses = (
      filtered_results.negative
       | "Generate Response" >> beam.Map(lambda x: ((x[0], x[3]), "<prompt>Generate an appology reponse for the user in this chat text: " + x[1] + "<answer>"))
       | "Gemma-Response" >> gemma_inference
       )

  generated_responses | "Print Reponse" >> beam.Map(lambda x: logging.info(x))
  filtered_results.main | "Print Main" >> beam.Map(lambda x: logging.info(x))
  filtered_results.error | "Print Errors" >> beam.Map(lambda x: logging.info(x))