from sagemaker.model import Model
from sagemaker import get_execution_role
from sagemaker.huggingface import get_huggingface_llm_image_uri
from sagemaker.huggingface import HuggingFaceModel
# retrieve the llm image uri
llm_image = get_huggingface_llm_image_uri(
  "huggingface",
  version="0.8.2"
)

role = get_execution_role()
hf_model_id = "tiiuae/falcon-7b-instruct" # model id from huggingface.co/models
model_name = hf_model_id.replace("/","-").replace(".","-")
endpoint_name = "sagemaker-fastchat-01"
instance_type = "ml.g5.2xlarge" # instance type to use for deployment
number_of_gpus = 1 # number of gpus to use for inference and tensor parallelism
health_check_timeout = 2400 # Increase the timeout for the health check to 5 minutes for downloading the model

llm_model = HuggingFaceModel(
      role=role,
      image_uri=llm_image,
      env={
        'HF_MODEL_ID': hf_model_id,
        # 'HF_MODEL_QUANTIZE': "bitsandbytes", # comment in to quantize
        'SM_NUM_GPUS': f"{number_of_gpus}",
        'MAX_INPUT_LENGTH': "3000",  # Max length of input text
        'MAX_TOTAL_TOKENS': "6000",  # Max length of the generation (including input text)
      },
      name=model_name
    )

llm = llm_model.deploy(
  initial_instance_count=1,
  instance_type=instance_type,
  container_startup_health_check_timeout=health_check_timeout,
  endpoint_name=endpoint_name,
)
sagemaker-inference/deploy_endpoint.py
aws-samples/genai-playground-common-design-patterns
"""
Chatbot Arena (side-by-side) tab.
Users chat with two chosen models.
"""

import dataclasses
import time
import json
from typing import List
import gradio as gr
import numpy as np
from io import StringIO
from concurrent.futures import ThreadPoolExecutor
from functools import partial

from fastchat.constants import (
    MODERATION_MSG,
    CONVERSATION_LIMIT_MSG,
    INACTIVE_MSG,
    INPUT_CHAR_LEN_LIMIT,
    CONVERSATION_TURN_LIMIT,
)
from fastchat.serve.gradio_web_server import (
    no_change_btn,
    disable_btn,
    get_model_description_md,
    ip_expiration_dict,
    State
)
from fastchat.utils import (
    build_logger,
    violates_moderation,
)

from fastchat.serve.gradio_block_arena_named import (
    regenerate,
    flash_buttons,
    clear_history
)
from fastchat.conversation import Conversation
from gradio_utils import (
    bot_response,
    upload_to_s3,
    get_conv_log_filename,
    conv_log
)


class SageMakerState(State):
    def __init__(self, model_name):
        super().__init__(model_name)
        self.conv.__class__ = SageMakerConversation
        self.conv.visible_messages = self.conv.messages.copy()

    def dict(self):
        base = self.conv.dict()
        base.update(
            {
                "conv_id": self.conv_id,
                "model_name": self.model_name,
            }
        )
        return base


@dataclasses.dataclass
class SageMakerConversation(Conversation):

    visible_messages: List[List[str]] = ()

    def append_message(self, role: str, message: str, visible_message: str = None):
        super().append_message(role, message)
        if visible_message:
            self.visible_messages.append([role, visible_message])
        else:
            self.visible_messages.append([role, message])

    def to_gradio_chatbot(self):
        """Convert the conversation to gradio chatbot format."""
        ret = []
        for i, (role, msg) in enumerate(self.visible_messages[self.offset:]):
            if i % 2 == 0:
                ret.append([msg, None])
            else:
                ret[-1][-1] = msg
        return ret

    def update_last_message(self, message: str):
        """Update the last output.

        The last message is typically set to be None when constructing the prompt,
        so we need to update it in-place after getting the response from a model.
        """
        super().update_last_message(message)
        self.visible_messages[-1][1] = message

    def copy(self):
        return SageMakerConversation(
            name=self.name,
            system_template=self.system_template,
            system_message=self.system_message,
            roles=self.roles,
            messages=[[x, y] for x, y in self.messages],
            visible_messages=[[x, y] for x, y in self.visible_messages],
            offset=self.offset,
            sep_style=self.sep_style,
            sep=self.sep,
            sep2=self.sep2,
            stop_str=self.stop_str,
            stop_token_ids=self.stop_token_ids,
        )

    def dict(self):
        ret = super().dict()
        ret.update({"visible_messages": self.visible_messages})
        return ret


logger = build_logger("gradio_web_server_multi", "gradio_web_server_multi.log")

num_sides = 2
enable_moderation = False


def vote_last_response(states, vote_type, model_selectors, request: gr.Request):
    filename = get_conv_log_filename(states[0].conv_id, "vote")
    data = {
        "tstamp": round(time.time(), 4),
        "type": vote_type,
        "models": [x for x in model_selectors],
        "states": [x.dict() for x in states],
        "ip": request.client.host,
    }
    upload_to_s3(
        conv_log.log_bucket,
        conv_log.log_prefix,
        filename,
        json.dumps(data)
        )


def leftvote_last_response(
    state0, state1, model_selector0, model_selector1, request: gr.Request
):
    logger.info(f"leftvote (named). ip: {request.client.host}")
    vote_last_response(
        [state0, state1], "leftvote", [model_selector0, model_selector1], request
    )
    return ("",) + (disable_btn,) * 4


def rightvote_last_response(
    state0, state1, model_selector0, model_selector1, request: gr.Request
):
    logger.info(f"rightvote (named). ip: {request.client.host}")
    vote_last_response(
        [state0, state1], "rightvote", [model_selector0, model_selector1], request
    )
    return ("",) + (disable_btn,) * 4


def tievote_last_response(
    state0, state1, model_selector0, model_selector1, request: gr.Request
):
    logger.info(f"tievote (named). ip: {request.client.host}")
    vote_last_response(
        [state0, state1], "tievote", [model_selector0, model_selector1], request
    )
    return ("",) + (disable_btn,) * 4


def bothbad_vote_last_response(
    state0, state1, model_selector0, model_selector1, request: gr.Request
):
    logger.info(f"bothbad_vote (named). ip: {request.client.host}")
    vote_last_response(
        [state0, state1], "bothbad_vote", [model_selector0, model_selector1], request
    )
    return ("",) + (disable_btn,) * 4


def load_webui_side_by_side_named(models, url_params):
    states = (None,) * num_sides

    model_left = models[0] if len(models) > 0 else ""
    if len(models) > 1:
        weights = ([8] * 4 + [4] * 8 + [1] * 32)[: len(models) - 1]
        weights = weights / np.sum(weights)
        model_right = np.random.choice(models[1:], p=weights)
    else:
        model_right = model_left

    selector_updates = (
        gr.Dropdown.update(choices=models, value=model_left, visible=True),
        gr.Dropdown.update(choices=models, value=model_right, visible=True),
    )

    return (
        states
        + selector_updates
        + (gr.Chatbot.update(visible=True),) * num_sides
        + (
            gr.Textbox.update(visible=True),
            gr.Box.update(visible=True),
            gr.Row.update(visible=True),
            gr.Row.update(visible=True),
            gr.Accordion.update(visible=True),
        )
    )


def add_text(
    state0, state1, model_selector0, model_selector1, text, request: gr.Request
):
    ip = request.client.host
    logger.info(f"add_text (named). ip: {ip}. len: {len(text)}")
    states = [state0, state1]
    model_selectors = [model_selector0, model_selector1]

    # Init states if necessary
    for i in range(num_sides):
        if states[i] is None:
            states[i] = SageMakerState(model_selectors[i])

    if len(text) <= 0:
        for i in range(num_sides):
            states[i].skip_next = True
        return (
            states
            + [x.to_gradio_chatbot() for x in states]
            + [""]
            + [
                no_change_btn,
            ]
            * 6
        )

    if ip_expiration_dict[ip] < time.time():
        logger.info(f"inactive (named). ip: {request.client.host}. text: {text}")
        for i in range(num_sides):
            states[i].skip_next = True
        return (
            states
            + [x.to_gradio_chatbot() for x in states]
            + [INACTIVE_MSG]
            + [
                no_change_btn,
            ]
            * 6
        )

    if enable_moderation:
        flagged = violates_moderation(text)
        if flagged:
            logger.info(
                f"violate moderation (named). ip: {request.client.host}. text: {text}"
            )
            for i in range(num_sides):
                states[i].skip_next = True
            return (
                states
                + [x.to_gradio_chatbot() for x in states]
                + [MODERATION_MSG]
                + [
                    no_change_btn,
                ]
                * 6
            )

    conv = states[0].conv
    if (len(conv.messages) - conv.offset) // 2 >= CONVERSATION_TURN_LIMIT:
        logger.info(f"conversation turn limit. ip: {request.client.host}. text: {text}")
        for i in range(num_sides):
            states[i].skip_next = True
        return (
            states
            + [x.to_gradio_chatbot() for x in states]
            + [CONVERSATION_LIMIT_MSG]
            + [
                no_change_btn,
            ]
            * 6
        )

    text = text[:INPUT_CHAR_LEN_LIMIT]  # Hard cut-off
    for i in range(num_sides):
        states[i].conv.append_message(states[i].conv.roles[0], text)
        states[i].conv.append_message(states[i].conv.roles[1], None)
        states[i].skip_next = False

    return (
        states
        + [x.to_gradio_chatbot() for x in states]
        + [""]
        + [
            disable_btn,
        ]
        * 6
    )


def send_file(state0, state1, model_selector0, model_selector1, file_box, request: gr.Request):
    logger.info(f"send_file. ip: {request.client.host}")
    states = [
        state0 if state0 is not None else SageMakerState(model_selector0),
        state1 if state1 is not None else SageMakerState(model_selector1)
        ]
    for state in states:
        if file_box is None:
            return states + [x.to_gradio_chatbot() for x in states] + [no_change_btn] * 6
        conv = state.conv
        if file_box.name.lower().endswith((".txt", ".csv", ".sh", ".py", ",java")):
            stringio = StringIO(open(file_box.name).read())
            content = "=== BEGIN FILE ===\n"
            content += stringio.read().strip()
            content += "\n=== END FILE ===\n\n Instruction: Please confirm that you have read that file by saying: 'Yes, I have read the file'"
            conv.append_message(conv.roles[0], content, "I have uploaded a file. Please confirm that you have read that file.")
            conv.append_message(conv.roles[1], None)
        else:
            conv.append_message(conv.roles[0], "Uploaded file")
            conv.append_message(conv.roles[1], None)
    return states + [x.to_gradio_chatbot() for x in states] + [disable_btn] * 6


def bot_response_multi(
    state0,
    state1,
    temperature,
    top_p,
    max_new_tokens,
    request: gr.Request,
):
    def collector_func(gen, future):
        gen.append(future.result())

    logger.info(f"bot_response_multi (named). ip: {request.client.host}")

    if state0.skip_next:
        # This generate call is skipped due to invalid inputs
        yield (
            state0,
            state1,
            state0.to_gradio_chatbot(),
            state1.to_gradio_chatbot(),
        ) + (no_change_btn,) * 6
        return

    states = [state0, state1]
    gen = []
    with ThreadPoolExecutor() as executor:

        for i in range(num_sides):
            kwargs = {
                "state": states[i],
                "temperature": temperature,
                "top_p": top_p,
                "max_new_tokens": max_new_tokens,
                "request": request
            }
            future = executor.submit(bot_response, **kwargs)
            future.add_done_callback(partial(collector_func, gen))

    chatbots = [None] * num_sides
    while True:
        stop = True
        for i in range(num_sides):
            try:
                ret = next(gen[i])
                states[i], chatbots[i] = ret[0], ret[1]
                stop = False
            except StopIteration:
                pass
        yield states + chatbots + [disable_btn] * 6
        if stop:
            break


def build_side_by_side_ui_named(models):
    notice_markdown = """
# ⚔️  Chatbot Arena ⚔️ 
### Rules
- Chat with two models side-by-side and vote for which one is better!
- You pick the models you want to chat with.
- You can do multiple rounds of conversations before voting.
- Click "Clear history" to start a new round.

### Choose two models to chat with
"""

    states = [gr.State() for _ in range(num_sides)]
    model_selectors = [None] * num_sides
    chatbots = [None] * num_sides

    with gr.Box(elem_id="share-region-named"):
        with gr.Row():
            for i in range(num_sides):
                with gr.Column():
                    model_selectors[i] = gr.Dropdown(
                        choices=models,
                        value=models[i] if len(models) > i else "",
                        interactive=True,
                        show_label=False,
                        container=False,
                    )

        with gr.Row():
            for i in range(num_sides):
                label = "Model A" if i == 0 else "Model B"
                with gr.Column():
                    chatbots[i] = gr.Chatbot(
                        label=label, elem_id="chatbot", visible=False, height=550
                    )

        with gr.Box() as button_row:
            with gr.Row():
                leftvote_btn = gr.Button(value="👈  A is better", interactive=False)
                rightvote_btn = gr.Button(value="👉  B is better", interactive=False)
                tie_btn = gr.Button(value="🤝  Tie", interactive=False)
                bothbad_btn = gr.Button(value="👎  Both are bad", interactive=False)

    with gr.Row():
        with gr.Column(scale=20):
            textbox = gr.Textbox(
                show_label=False,
                placeholder="Enter text and press ENTER",
                visible=False,
                container=False,
            )
        with gr.Column(scale=1, min_width=50):
            send_file_btn = gr.UploadButton(label="📂️", visible=False, size="lg", interactive=True, type="file", file_types=["text"])

    with gr.Row() as button_row2:
        regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
        clear_btn = gr.Button(value="🗑️  Clear history", interactive=False)

    with gr.Accordion("Parameters", open=False, visible=True) as parameter_row:
        temperature = gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.7,
            step=0.1,
            interactive=True,
            label="Temperature",
        )
        top_p = gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=1.0,
            step=0.1,
            interactive=True,
            label="Top P",
        )
        max_output_tokens = gr.Slider(
            minimum=16,
            maximum=4096,
            value=512,
            step=64,
            interactive=True,
            label="Max output tokens",
        )

    # Register listeners
    btn_list = [
        leftvote_btn,
        rightvote_btn,
        tie_btn,
        bothbad_btn,
        regenerate_btn,
        clear_btn,
    ]
    leftvote_btn.click(
        leftvote_last_response,
        states + model_selectors,
        [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
    )
    rightvote_btn.click(
        rightvote_last_response,
        states + model_selectors,
        [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
    )
    tie_btn.click(
        tievote_last_response,
        states + model_selectors,
        [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
    )
    bothbad_btn.click(
        bothbad_vote_last_response,
        states + model_selectors,
        [textbox, leftvote_btn, rightvote_btn, tie_btn, bothbad_btn],
    )
    regenerate_btn.click(
        regenerate, states, states + chatbots + [textbox] + btn_list
    ).then(
        bot_response_multi,
        states + [temperature, top_p, max_output_tokens],
        states + chatbots + btn_list,
    ).then(
        flash_buttons, [], btn_list
    )
    clear_btn.click(clear_history, None, states + chatbots + [textbox] + btn_list)

    for i in range(num_sides):
        model_selectors[i].change(
            clear_history, None, states + chatbots + [textbox] + btn_list
        )

    textbox.submit(
        add_text,
        states + model_selectors + [textbox],
        states + chatbots + [textbox] + btn_list,
    ).then(
        bot_response_multi,
        states + [temperature, top_p, max_output_tokens],
        states + chatbots + btn_list,
    ).then(
        flash_buttons, [], btn_list
    )
    send_file_btn.upload(
        send_file, states + model_selectors + [send_file_btn], states + chatbots + btn_list
    ).then(
        bot_response_multi,
        states + [temperature, top_p, max_output_tokens],
        states + chatbots + btn_list,
    )

    return (
        states,
        model_selectors,
        chatbots,
        textbox,
        send_file_btn,
        button_row,
        button_row2,
        parameter_row,
    )
serve/gradio_block_arena_named.py
aws-samples/genai-playground-common-design-patterns
import os
import requests
import datetime
import time
import json
import gradio as gr
import boto3

from fastchat.serve.gradio_web_server import (
    logger,
    enable_btn,
    disable_btn,
    no_change_btn,
    model_worker_stream_iter,
    post_process_code
)
from fastchat.constants import (
    SERVER_ERROR_MSG,
    ErrorCode
)
from fastchat.serve.api_provider import (
    anthropic_api_stream_iter,
    openai_api_stream_iter,
    palm_api_stream_iter
)

worker_addr = "http://127.0.0.1:21002"

class conv_logger:
    def __init__(self):
        self.log_bucket = ""
        self.log_prefix = ""


def upload_to_s3(bucket, prefix, filename, data):
    s3_client.put_object(
        Body=data, 
        Bucket=bucket, 
        Key=os.path.join(prefix, filename)
    )


def get_conv_log_filename(conv_id, log_type):
    t = datetime.datetime.utcnow()
    name = f"conv-{t.hour:02d}-{t.minute:02d}-{t.second:02d}.json"
    prefix = f"{log_type}/{t.year}/{t.month:02d}/{t.day:02d}/{conv_id}"
    return os.path.join(prefix, name)


def bot_response(state, temperature, top_p, max_new_tokens, request: gr.Request):
    logger.info(f"bot_response. ip: {request.client.host}")
    start_tstamp = time.time()
    temperature = float(temperature)
    top_p = float(top_p)
    max_new_tokens = int(max_new_tokens)

    if state.skip_next:
        # This generate call is skipped due to invalid inputs
        state.skip_next = False
        yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5
        return

    conv, model_name = state.conv, state.model_name
    if model_name == "gpt-3.5-turbo" or model_name == "gpt-4":
        prompt = conv.to_openai_api_messages()
        stream_iter = openai_api_stream_iter(
            model_name, prompt, temperature, top_p, max_new_tokens
        )
    elif model_name == "claude-2" or model_name == "claude-instant-1":
        prompt = conv.get_prompt()
        stream_iter = anthropic_api_stream_iter(
            model_name, prompt, temperature, top_p, max_new_tokens
        )
    elif model_name == "palm-2":
        stream_iter = palm_api_stream_iter(
            state.palm_chat, conv.messages[-2][1], temperature, top_p, max_new_tokens
        )
    else:
        logger.info(f"model_name: {model_name}, worker_addr: {worker_addr}")

        # No available worker
        if worker_addr == "":
            conv.update_last_message(SERVER_ERROR_MSG)
            yield (
                state,
                state.to_gradio_chatbot(),
                disable_btn,
                disable_btn,
                disable_btn,
                enable_btn,
                enable_btn,
            )
            return

        # Construct prompt.
        # We need to call it here, so it will not be affected by "▌".
        prompt = conv.get_prompt()

        # Set repetition_penalty
        if "t5" in model_name:
            repetition_penalty = 1.2
        else:
            repetition_penalty = 1.0

        stream_iter = model_worker_stream_iter(
            conv,
            model_name,
            worker_addr,
            prompt,
            temperature,
            repetition_penalty,
            top_p,
            max_new_tokens,
        )

    conv.update_last_message("▌")
    yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5

    try:
        for data in stream_iter:
            if data["error_code"] == 0:
                output = data["text"].strip()
                if "vicuna" in model_name:
                    output = post_process_code(output)
                conv.update_last_message(output + "▌")
                yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
            else:
                output = data["text"] + f"\n\n(error_code: {data['error_code']})"
                conv.update_last_message(output)
                yield (state, state.to_gradio_chatbot()) + (
                    disable_btn,
                    disable_btn,
                    disable_btn,
                    enable_btn,
                    enable_btn,
                )
                return
            time.sleep(0.015)
    except requests.exceptions.RequestException as e:
        conv.update_last_message(
            f"{SERVER_ERROR_MSG}\n\n"
            f"(error_code: {ErrorCode.GRADIO_REQUEST_ERROR}, {e})"
        )
        yield (state, state.to_gradio_chatbot()) + (
            disable_btn,
            disable_btn,
            disable_btn,
            enable_btn,
            enable_btn,
        )
        return
    except Exception as e:
        conv.update_last_message(
            f"{SERVER_ERROR_MSG}\n\n"
            f"(error_code: {ErrorCode.GRADIO_STREAM_UNKNOWN_ERROR}, {e})"
        )
        yield (state, state.to_gradio_chatbot()) + (
            disable_btn,
            disable_btn,
            disable_btn,
            enable_btn,
            enable_btn,
        )
        return

    # Delete "▌"
    conv.update_last_message(conv.messages[-1][-1][:-1])
    yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 5

    finish_tstamp = time.time()
    logger.info(f"{output}")

    filename = get_conv_log_filename(state.conv_id, "conv")
    data = {
        "tstamp": round(finish_tstamp, 4),
        "type": "chat",
        "model": model_name,
        "gen_params": {
            "temperature": temperature,
            "top_p": top_p,
            "max_new_tokens": max_new_tokens,
        },
        "start": round(start_tstamp, 4),
        "finish": round(finish_tstamp, 4),
        "state": state.dict(),
        "ip": request.client.host,
    }
    upload_to_s3(
        conv_log.log_bucket,
        conv_log.log_prefix,
        filename,
        json.dumps(data)
        )


s3_client = boto3.client("s3")
conv_log = conv_logger()
serve/gradio_utils.py
aws-samples/genai-playground-common-design-patterns
import json
from jsonpath_ng import jsonpath, parse

class BaseModel:
    def inovke_api(self):
        raise NotImplementedError("Not implemented in base model, make sure to override this method.")
serve/handlers/base.py
aws-samples/genai-playground-common-design-patterns
import os
import io
import json
import boto3
from handlers.base import BaseModel


class model(BaseModel):
    def __init__(self, model_name):
        super().__init__()
        self.model_name = model_name
        self.bedrock_client = boto3.client(
            service_name="bedrock-runtime", 
            region_name="us-west-2"
        )
        self.invoke_api = self.bedrock_client.invoke_model_with_response_stream

    def invoke(self, body):
        try:
            response = self.invoke_api(
                modelId=self.model_name,
                body = json.dumps(body).encode("utf-8")
            )["body"]
            return response
        except Exception as e:
            print(f"Error {e}, Body {body}")
            raise e
            
class StreamIterator:
    def __init__(self, stream):
        self.byte_iterator = iter(stream)
        self.buffer = io.BytesIO()
        self.read_pos = 0

    def __iter__(self):
        return self

    def __next__(self):
        while True:
            self.buffer.seek(self.read_pos)
            line = self.buffer.readline()
            if line:
                self.read_pos += len(line)
                return line
            try:
                chunk = next(self.byte_iterator)
            except StopIteration:
                if self.read_pos < self.buffer.getbuffer().nbytes:
                    continue
                raise
            if 'chunk' not in chunk:
                print(f"Unknown event type: {chunk}")
                continue
            self.buffer.seek(0, io.SEEK_END)
            self.buffer.write(chunk['chunk']['bytes'])
        

serve/handlers/bedrock.py
aws-samples/genai-playground-common-design-patterns
import os
import io
import json
import boto3
from handlers.base import BaseModel

class model(BaseModel):
    def __init__(self, model_name):
        super().__init__()
        self.endpoint_name = model_name
        self.sagemaker_client = boto3.client(
            service_name="sagemaker-runtime"
        )
        self.invoke_api = self.sagemaker_client.invoke_endpoint_with_response_stream

    def invoke(self, body):
        response = self.invoke_api(
            EndpointName=self.endpoint_name,
            Body = json.dumps(body).encode("utf-8"),
            ContentType="application/json"
        )["Body"]
        return response
    
class StreamIterator:
    def __init__(self, stream):
        self.byte_iterator = iter(stream)
        self.buffer = io.BytesIO()
        self.read_pos = 0

    def __iter__(self):
        return self

    def __next__(self):
        while True:
            self.buffer.seek(self.read_pos)
            line = self.buffer.readline()
            if line and line[-1] == 10:
                self.read_pos += len(line)
                return line[:-1]
            try:
                chunk = next(self.byte_iterator)
            except StopIteration:
                if self.read_pos < self.buffer.getbuffer().nbytes:
                    continue
                raise
            if 'PayloadPart' not in chunk:
                print(f"Unknown event type: {chunk}")
                continue
            self.buffer.seek(0, io.SEEK_END)
            self.buffer.write(chunk['PayloadPart']['Bytes'])
serve/handlers/sagemaker.py
aws-samples/genai-playground-common-design-patterns
"""
The gradio webui server for chatting with a single model.
"""
import argparse
import time
import json
from urllib.parse import urlparse
import gradio as gr
import boto3
from fastchat.utils import (
    get_window_url_params_js,
    parse_gradio_auth_creds,
    violates_moderation
)
from fastchat.serve.gradio_web_server import (
    logger,
    enable_moderation,
    block_css,
    flag_last_response,
    ip_expiration_dict,
    disable_btn,
    no_change_btn,
)
from fastchat.serve.gradio_block_arena_named import (
    set_global_vars_named,
)
from fastchat.constants import (
    SESSION_EXPIRATION_TIME,
    CONVERSATION_LIMIT_MSG,
    CONVERSATION_TURN_LIMIT,
    INACTIVE_MSG,
    MODERATION_MSG,
    INPUT_CHAR_LEN_LIMIT,
)
from fastchat.model.model_adapter import (
    BaseModelAdapter,
    model_adapters
)
from fastchat.conversation import Conversation, get_conv_template
import fastchat
from gradio_block_arena_named import (
    build_side_by_side_ui_named,
    load_webui_side_by_side_named,
    SageMakerState
)
from gradio_utils import (
    bot_response,
    upload_to_s3,
    get_conv_log_filename,
    conv_log,
    worker_addr
)
import gradio_utils as utils
from io import StringIO

s3_client = boto3.client("s3")
ddb_client = boto3.client("dynamodb")
conv_templates = ""


def set_global_vars(ddb_table_name_, enable_moderation_, conv_templates_):
    global ddb_table_name, enable_moderation, conv_templates
    ddb_table_name = ddb_table_name_
    enable_moderation = enable_moderation_
    conv_templates = conv_templates_


def upload_file(files):
    file_paths = [file.name for file in files]
    return file_paths


def load_webui_single(models, url_params):
    selected_model = models[0] if len(models) > 0 else ""
    if "model" in url_params:
        model = url_params["model"]
        if model in models:
            selected_model = model

    dropdown_update = gr.Dropdown.update(
        choices=models, value=selected_model, visible=True
    )

    state = None
    return (
        state,
        dropdown_update,
        gr.Chatbot.update(visible=True),
        gr.Textbox.update(visible=True),
        gr.UploadButton.update(visible=True),
        gr.Row.update(visible=True),
        gr.Accordion.update(visible=True),
    )


def add_text(state, model_selector, text, request: gr.Request):
    ip = request.client.host
    logger.info(f"add_text. ip: {ip}. len: {len(text)}")

    if state is None:
        state = SageMakerState(model_selector)

    if len(text) <= 0:
        state.skip_next = True
        return (state, state.to_gradio_chatbot(), "") + (no_change_btn,) * 5

    if ip_expiration_dict[ip] < time.time():
        logger.info(f"inactive. ip: {request.client.host}. text: {text}")
        state.skip_next = True
        return (state, state.to_gradio_chatbot(), INACTIVE_MSG) + (no_change_btn,) * 5

    if enable_moderation:
        flagged = violates_moderation(text)
        if flagged:
            logger.info(f"violate moderation. ip: {request.client.host}. text: {text}")
            state.skip_next = True
            return (state, state.to_gradio_chatbot(), MODERATION_MSG) + (
                no_change_btn,
            ) * 5

    conv = state.conv
    if (len(conv.messages) - conv.offset) // 2 >= CONVERSATION_TURN_LIMIT:
        logger.info(f"conversation turn limit. ip: {request.client.host}. text: {text}")
        state.skip_next = True
        return (state, state.to_gradio_chatbot(), CONVERSATION_LIMIT_MSG) + (
            no_change_btn,
        ) * 5

    text = text[:INPUT_CHAR_LEN_LIMIT]  # Hard cut-off
    conv.append_message(conv.roles[0], text)
    conv.append_message(conv.roles[1], None)
    return (state, state.to_gradio_chatbot(), "") + (disable_btn,) * 5


def load_webui(url_params, request: gr.Request):
    global models

    ip = request.client.host
    logger.info(f"load_webui. ip: {ip}. params: {url_params}")
    ip_expiration_dict[ip] = time.time() + SESSION_EXPIRATION_TIME

    if args.model_list_mode == "reload":
        models = get_model_list(
            ddb_table_name
        )

    return load_webui_single(models, url_params)


def send_file(state, model_selector, file_box, request: gr.Request):
    logger.info(f"send_file. ip: {request.client.host}")
    if state is None:
        state = SageMakerState(model_selector)

    if file_box is None:
        return (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5

    conv = state.conv

    if file_box.name.lower().endswith((".txt", ".csv", ".sh", ".py", ",java")):
        stringio = StringIO(open(file_box.name).read())
        content = "=== BEGIN FILE ===\n"
        content += stringio.read().strip()
        content += "\n=== END FILE ===\n\n Instruction: Please confirm that you have read that file by saying: 'Yes, I have read the file'"

        conv.append_message(conv.roles[0], content, "I have uploaded a file. Please confirm that you have read that file.")
        conv.append_message(conv.roles[1], None)
    else:
        conv.append_message(conv.roles[0], "Uploaded file")
        conv.append_message(conv.roles[1], None)

    return (state, state.to_gradio_chatbot()) + (disable_btn,) * 5


def regenerate(state, request: gr.Request):
    logger.info(f"regenerate. ip: {request.client.host}")
    state.conv.update_last_message(None)
    load_templates_from_s3(conv_templates)
    return (state, state.to_gradio_chatbot(), "") + (disable_btn,) * 5


def clear_history(request: gr.Request):
    logger.info(f"clear_history. ip: {request.client.host}")
    state = None
    load_templates_from_s3(conv_templates)
    ip = request.client.host
    ip_expiration_dict[ip] = time.time() + SESSION_EXPIRATION_TIME
    return (state, [], "") + (disable_btn,) * 5


def build_single_model_ui(models):

    notice_markdown = f"""
# 🏔️ Chat with Open Large Language Models

### Choose a model to chat with
"""

    state = gr.State()
    gr.Markdown(notice_markdown, elem_id="notice_markdown")

    with gr.Row(elem_id="model_selector_row"):
        model_selector = gr.Dropdown(
            choices=models,
            value=models[0] if len(models) > 0 else "",
            interactive=True,
            show_label=False,
            container=False,
        )

    chatbot = gr.Chatbot(
        elem_id="chatbot",
        label="Scroll down and start chatting",
        visible=False,
        height=550,
    )
    with gr.Row():
        with gr.Column(scale=20):
            textbox = gr.Textbox(
                show_label=False,
                placeholder="Enter text and press ENTER",
                visible=False,
                container=False,
            )
        with gr.Column(scale=1, min_width=50):
            send_file_btn = gr.UploadButton(label="📂️", visible=False, size="lg", interactive=True, type="file", file_types=["text"])

    with gr.Row(visible=False) as button_row:
        upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
        downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
        flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
        regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
        clear_btn = gr.Button(value="🗑️  Clear history", interactive=False)

    with gr.Accordion("Parameters", open=False, visible=False) as parameter_row:
        temperature = gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.7,
            step=0.1,
            interactive=True,
            label="Temperature",
        )
        top_p = gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=1.0,
            step=0.1,
            interactive=True,
            label="Top P",
        )
        max_output_tokens = gr.Slider(
            minimum=16,
            maximum=8196,
            value=512,
            step=64,
            interactive=True,
            label="Max output tokens",
        )

    # Register listeners
    btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
    upvote_btn.click(
        upvote_last_response,
        [state, model_selector],
        [textbox, upvote_btn, downvote_btn, flag_btn],
    )
    downvote_btn.click(
        downvote_last_response,
        [state, model_selector],
        [textbox, upvote_btn, downvote_btn, flag_btn],
    )
    flag_btn.click(
        flag_last_response,
        [state, model_selector],
        [textbox, upvote_btn, downvote_btn, flag_btn],
    )
    regenerate_btn.click(regenerate, state, [state, chatbot, textbox] + btn_list).then(
        bot_response,
        [state, temperature, top_p, max_output_tokens],
        [state, chatbot] + btn_list,
    )
    clear_btn.click(clear_history, None, [state, chatbot, textbox] + btn_list)

    model_selector.change(clear_history, None, [state, chatbot, textbox] + btn_list)

    textbox.submit(
        add_text, [state, model_selector, textbox], [state, chatbot, textbox] + btn_list
    ).then(
        bot_response,
        [state, temperature, top_p, max_output_tokens],
        [state, chatbot] + btn_list,
    )
    send_file_btn.upload(
        send_file, [state, model_selector, send_file_btn], [state, chatbot] + btn_list
    ).then(
        bot_response,
        [state, temperature, top_p, max_output_tokens],
        [state, chatbot] + btn_list,
    )

    return (
        state, 
        model_selector, 
        chatbot,
        textbox, 
        send_file_btn, 
        button_row, 
        parameter_row
        )


def build_webui_single(models):
    with gr.Blocks(
        title="Chat with Open Large Language Models hosted on SageMaker",
        theme=gr.themes.Base(),
        css=block_css,
    ) as webui:
        url_params = gr.JSON(visible=False)

        (
            state,
            model_selector,
            chatbot,
            textbox, 
            send_file_btn,
            button_row,
            parameter_row
        ) = build_single_model_ui(models)

        if args.model_list_mode not in ["once", "reload"]:
            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
        webui.load(
            load_webui,
            [url_params],
            [
                state,
                model_selector,
                chatbot,
                textbox,
                send_file_btn,
                button_row,
                parameter_row
            ],
            _js=get_window_url_params_js,
        )

    return webui


def load_webui_multi(url_params, request: gr.Request):
    global models

    ip = request.client.host
    logger.info(f"load_webui_multi. ip: {ip}. params: {url_params}")
    ip_expiration_dict[ip] = time.time() + SESSION_EXPIRATION_TIME

    selected = 0
    if "compare" in url_params:
        selected = 1

    models = get_model_list(
        ddb_table_name
    )

    single_updates = load_webui_single(models, url_params)

    side_by_side_named_updates = load_webui_side_by_side_named(models, url_params)
    return (
        (gr.Tabs.update(selected=selected),)
        + single_updates
        + side_by_side_named_updates
    )


def build_webui_multi(models):
    with gr.Blocks(
        title="Chat with Open Large Language Models",
        theme=gr.themes.Base(),
        css=block_css,
    ) as webui:
        with gr.Tabs() as tabs:
            with gr.Tab("Single Model", id=0):
                (
                    a_state,
                    a_model_selector,
                    a_chatbot,
                    a_textbox,
                    a_send_btn,
                    a_button_row,
                    a_parameter_row,
                ) = build_single_model_ui(models)
                a_list = [
                    a_state,
                    a_model_selector,
                    a_chatbot,
                    a_textbox,
                    a_send_btn,
                    a_button_row,
                    a_parameter_row,
                ]

            with gr.Tab("Chatbot Arena (side-by-side)", id=1):
                (
                    b_states,
                    b_model_selectors,
                    b_chatbots,
                    b_textbox,
                    b_send_btn,
                    b_button_row,
                    b_button_row2,
                    b_parameter_row,
                ) = build_side_by_side_ui_named(models)
                b_list = (
                    b_states
                    + b_model_selectors
                    + b_chatbots
                    + [
                        b_textbox,
                        b_send_btn,
                        b_button_row,
                        b_button_row2,
                        b_parameter_row,
                    ]
                )

        url_params = gr.JSON(visible=False)

        if args.model_list_mode not in ["once", "reload"]:
            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
        webui.load(
            load_webui_multi,
            [url_params],
            [tabs] + a_list + b_list,
            _js=get_window_url_params_js,
        )

    return webui


def load_templates_from_s3(url):
    if url:
        scheme, netloc, path, _, _, _ = urlparse(url)

        if scheme == "s3":
            path = path[1:] if len(path) > 0 and path.startswith("/") else path
            if path.endswith(".json"):
                if netloc:
                    logger.info(f"load_templates: {url}")
                    result = s3_client.get_object(Bucket=netloc, Key=path)
                    json_templates = json.loads(
                        result["Body"].read().decode(),
                        object_hook=lambda d: Conversation(**d) if not isinstance(list(d.values())[0], Conversation) else d, 
                        strict=False
                        )
                    for conv_template, template in json_templates.items():
                        fastchat.conversation.conv_templates[conv_template] = template


def vote_last_response(state, vote_type, model_selector, request: gr.Request):
    filename = get_conv_log_filename(state.conv_id, "vote")
    data = {
        "tstamp": round(time.time(), 4),
        "type": vote_type,
        "model": model_selector,
        "state": state.dict(),
        "ip": request.client.host,
    }
    upload_to_s3(
        conv_log.log_bucket,
        conv_log.log_prefix,
        filename,
        json.dumps(data)
        )


def upvote_last_response(state, model_selector, request: gr.Request):
    logger.info(f"upvote. ip: {request.client.host}")
    vote_last_response(state, "upvote", model_selector, request)
    return ("",) + (disable_btn,) * 3


def downvote_last_response(state, model_selector, request: gr.Request):
    logger.info(f"downvote. ip: {request.client.host}")
    vote_last_response(state, "downvote", model_selector, request)
    return ("",) + (disable_btn,) * 3


def flag_last_response(state, model_selector, request: gr.Request):
    logger.info(f"flag. ip: {request.client.host}")
    vote_last_response(state, "flag", model_selector, request)
    return ("",) + (disable_btn,) * 3

def get_model_list(ddb_table_name):
    ddb_client = boto3.client("dynamodb")
    return ddb_client.get_item(
        TableName=ddb_table_name, 
        Key={
            "model_name": {
                "S": "all-models"
            }
        }
    )["Item"]["model_names"]["S"].split(",")

class ClaudeAdapter(BaseModelAdapter):
    """The model adapter for Claude"""

    def match(self, model_path: str):
        return model_path in ["anthropic.claude-v2", "anthropic.claude-instant-v1", "anthropic.claude-v1"]

    def load_model(self, model_path: str, from_pretrained_kwargs: dict):
        raise NotImplementedError()

    def get_default_conv_template(self, model_path: str) -> Conversation:
        return get_conv_template("claude")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int)
    parser.add_argument(
        "--share",
        action="store_true",
        help="Whether to generate a public, shareable link.",
    )
    parser.add_argument(
        "--concurrency-count",
        type=int,
        default=10,
        help="The concurrency count of the gradio queue.",
    )
    parser.add_argument(
        "--model-list-mode",
        type=str,
        default="once",
        choices=["once", "reload"],
        help="Whether to load the model list once or reload the model list every time.",
    )
    parser.add_argument(
        "--mode",
        type=str,
        default="single",
        choices=["single", "multi"],
        help="WebUI mode, single model or side by side multi",
    )
    parser.add_argument(
        "--moderate", action="store_true", help="Enable content moderation"
    )
    parser.add_argument(
        "--gradio-auth-path",
        type=str,
        help='Set the gradio authentication file path. The file should contain one or more user:password pairs in this format: "u1:p1,u2:p2,u3:p3"',
        default=None,
    )
    parser.add_argument(
        "--conv-templates", type=str, default=None, help="Conversation prompt template."
    )
    parser.add_argument(
        "--conv-bucket", type=str, default=None, help="S3 Bucket to store conversation logs"
    )
    parser.add_argument(
        "--conv-prefix", type=str, default=None, help="S3 prefix to store conversation"
    )
    parser.add_argument(
        "--ddb-table-name", type=str, default=None, help="DDB table for list of models"
    )
    parser.add_argument(
        "--worker-address", type=str, default="http://127.0.0.1:21002", help="Model worker address"
    )
    args = parser.parse_args()
    fastchat.serve.gradio_web_server.args = args
    logger.info(f"args: {args}")
    worker_addr = args.worker_address

    load_templates_from_s3(args.conv_templates)
    fastchat.model.model_adapter.model_adapters.append(ClaudeAdapter())

    # Set global variables
    set_global_vars(
        args.ddb_table_name,
        args.moderate,
        args.conv_templates
        )
    set_global_vars_named(args.moderate)
    conv_log.log_bucket = args.conv_bucket
    conv_log.log_prefix = args.conv_prefix
    fastchat.serve.gradio_block_arena_named.set_global_vars_named(args.moderate)
    models = get_model_list(
        args.ddb_table_name
    )

    # Set authorization credentials
    auth = None
    if args.gradio_auth_path is not None:
        auth = parse_gradio_auth_creds(args.gradio_auth_path)

    # Launch the webui
    if args.mode == "single":
        webui = build_webui_single(models)
    else:
        # Multi agent UI
        webui = build_webui_multi(models)
    webui.queue(
        concurrency_count=args.concurrency_count, status_update_rate=10, api_open=False
    ).launch(
        server_name=args.host,
        server_port=args.port,
        share=args.share,
        max_threads=200,
        auth=auth,
    )
serve/sagemaker_gradio_web_server.py
aws-samples/genai-playground-common-design-patterns
"""
A model worker that executes the model.
"""
import argparse
import json
from urllib.parse import urlparse
from jsonpath_ng import jsonpath, parse
import re
from typing import List, Dict
import urllib3
import io
import boto3
from botocore.response import StreamingBody
from botocore.eventstream import EventStream
from fastapi import Request
from fastapi.responses import StreamingResponse, JSONResponse
import uvicorn
from fastchat.constants import ErrorCode, SERVER_ERROR_MSG
from fastchat.model.model_adapter import add_model_args
from fastchat.serve.model_worker import (
    worker_id,
    logger,
    app,
    BaseModelWorker,
    release_worker_semaphore,
    acquire_worker_semaphore,
    create_background_tasks
)
import fastchat
from importlib import import_module


class ModelWorker(BaseModelWorker):
    def __init__(
        self,
        ddb_table_name: str,
        worker_id: str,
        model_path: str,
        limit_worker_concurrency: int,
        no_register: bool,
        conv_template: str = None
    ):
        self.ddb_table_name = ddb_table_name
        self.ddb_client = boto3.client("dynamodb")
        self.model_names = self.get_model_list()
        super().__init__(
            "",
            "",
            worker_id,
            model_path,
            self.model_names,
            limit_worker_concurrency,
            conv_template
        )

        logger.info("Starting model worker...")
        logger.info(f"Loading the model {self.model_names} on worker {worker_id} ...")
        self.http = urllib3.PoolManager()
        self.serving_port = 8080
        
    def get_model_list(self):
        return self.ddb_client.get_item(
            TableName=self.ddb_table_name, 
            Key={
                "model_name": {
                    "S": "all-models"
                }
            }
        )["Item"]["model_names"]["S"].split(",")

    def invoke_model(self, params, request_form):
        if params["top_p"] == 0:
            params["top_p"] = 0.01
        elif params["top_p"] == 1.0:
            params["top_p"] = 0.99
        if "stop" in params and params["stop"]:
            params["stop"] = [params["stop"]] if isinstance(params["stop"], str) else params["stop"]

        body = self.form_request(
            params,
            request_form["defaults"],
            request_form["mapping"]
        )

        response = self.invoke(
            body=body
        )
        if isinstance(response, urllib3.response.HTTPResponse):
            if response.getheader("content-type") == "text/event-stream":
                output = response.stream()
        elif isinstance(response, StreamingBody) or isinstance(response, EventStream):
            output = response
        else:
            output = response.read()
        return output

    def form_request(self, params, defaults, mapping):
        for attrib, jpath in mapping.items():
            if attrib in params.keys():
                jsonpath_expr = parse(jpath)
                jsonpath_expr.update(defaults, params[attrib])
        return defaults

    def parse_response(self, response_body, mapping, regex_sub=None):
        res = None
        if regex_sub:
            res = json.loads(re.sub(regex_sub, "", response_body))
        else:
            res = json.loads(response_body)
        ret = {}
        for attrib, jpath in mapping.items():
            jsonpath_expr = parse(jpath)
            results = jsonpath_expr.find(res)
            if results and len(results) > 0:
                ret[attrib] = results[0].value
        return ret
    
    def get_model_config(self, model_name):
        resp = self.ddb_client.get_item(
            TableName=self.ddb_table_name, 
            Key={
                "model_name": {
                "S": model_name
                }
            }
        )["Item"]
        json_config = json.loads(
            resp["endpoint_req_config"]["S"],
            strict=False
            )
        if "response" not in json_config.keys() or "request" not in json_config.keys():
            raise ValueError("Worker config file must container 'request' and 'response' keys.")
        model_family = resp["model_family"]["S"]
        if model_family == "sagemaker":
            model_name_ =  resp["endpoint_name"]["S"]
        else:
            model_name_ = model_name        
        self.invoke = import_module("handlers." + model_family).model(model_name_).invoke
        self.stream_iterator = import_module("handlers." + model_family).StreamIterator
        return (
            json_config["request"],
            json_config["response"]
        )

    def generate_stream_gate(self, params):
        self.call_ct += 1
        model_name = params["model"]
        (
            request_form,
            response_form
        ) = self.get_model_config(model_name)

        try:
            text = ""
            for line in self.stream_iterator(self.invoke_model(params, request_form)):
                if line:
                    output = self.parse_response(
                        line.decode("utf-8"),
                        response_form["mapping"],
                        response_form["regex_sub"]
                    )
                    if "error" in output:
                        raise ValueError(output)
                    if "stop" in params:
                        stop_token = params["stop"] if isinstance(params["stop"], list) else list(params["stop"])
                    else:
                        stop_token = None
                    if output["text"] not in stop_token:
                        text += output["text"]
                    else:
                        break
                    output["text"] = text
                output["error_code"] = 0
                yield json.dumps(output).encode() + b"\0"
        except (ValueError, RuntimeError) as e:
            ret = {
                "text": f"{SERVER_ERROR_MSG}\n\n({e})",
                "error_code": ErrorCode.INTERNAL_ERROR,
            }
            yield json.dumps(ret).encode() + b"\0"

    def generate_gate(self, params):
        for x in self.generate_stream_gate(params):
            pass
        return json.loads(x[:-1].decode())


@app.post("/worker_generate_stream")
async def api_generate_stream(request: Request):
    params = await request.json()
    await acquire_worker_semaphore()
    generator = worker.generate_stream_gate(params)
    background_tasks = create_background_tasks()
    return StreamingResponse(generator, background=background_tasks)


@app.post("/worker_generate")
async def api_generate(request: Request):
    params = await request.json()
    await acquire_worker_semaphore()
    output = worker.generate_gate(params)
    release_worker_semaphore()
    return JSONResponse(output)


@app.post("/worker_get_status")
async def api_get_status(request: Request):
    return worker.get_status()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="localhost")
    parser.add_argument("--port", type=int, default=21002)
    add_model_args(parser)
    parser.add_argument(
        "--conv-template", type=str, default=None, help="Conversation prompt template."
    )
    parser.add_argument(
        "--limit-worker-concurrency",
        type=int,
        default=5,
        help="Limit the model concurrency to prevent OOM.",
    )
    parser.add_argument("--stream-interval", type=int, default=2)
    parser.add_argument("--no-register", action="store_true")
    parser.add_argument("--ddb-table-name", type=str, default="")
    args = parser.parse_args()
    logger.info(f"args: {args}")

    worker = ModelWorker(
        args.ddb_table_name,
        worker_id,
        args.model_path,
        args.limit_worker_concurrency,
        no_register=args.no_register,
        conv_template=args.conv_template
    )

    fastchat.serve.model_worker.worker = worker
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")
serve/sagemaker_model_worker.py
aws-samples/genai-playground-common-design-patterns