UI and API implementation for base64 encoded image input

Author: vmpuri

torchchat/usages/browser.py

@@ -1,39 +1,123 @@
+import base64
+import logging
 import time
+from pathlib import Path
+
 import streamlit as st
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
 from openai import OpenAI
 
 st.title("torchchat")
 
 start_state = [
     {
         "role": "system",
-        "content": "You're an assistant. Answer questions directly, be brief, and have fun.",
+        "content": "You're a helpful assistant - have fun.",
     },
     {"role": "assistant", "content": "How can I help you?"},
 ]
 
+st.session_state.uploader_key = 0
+
+
+def reset_per_message_state():
+    # Catch all function for anything that should be reset between each message.
+    _update_uploader_key()
+
+
+def _update_uploader_key():
+    # Increment the uploader key to reset the file uploader after each message.
+    st.session_state.uploader_key = int(time.time())
+
+
 with st.sidebar:
+    # API Configuration
+    api_base_url = st.text_input(
+        label="API Base URL",
+        value="http://127.0.0.1:5000/v1",
+        help="The base URL for the OpenAI API to connect to",
+    )
+
+    st.divider()
+    temperature = st.slider(
+        "Temperature", min_value=0.0, max_value=1.0, value=1.0, step=0.01
+    )
+
     response_max_tokens = st.slider(
         "Max Response Tokens", min_value=10, max_value=1000, value=250, step=10
     )
     if st.button("Reset Chat", type="primary"):
         st.session_state["messages"] = start_state
 
+    image_prompts = st.file_uploader(
+        "Image Prompts",
+        type=["jpeg"],
+        accept_multiple_files=True,
+        key=st.session_state.uploader_key,
+    )
+
+    for image in image_prompts:
+        st.image(image)
+
+
+client = OpenAI(
+    base_url=api_base_url,
+    api_key="813",  # The OpenAI API requires an API key, but since we don't consume it, this can be any non-empty string.
+)
+
 if "messages" not in st.session_state:
     st.session_state["messages"] = start_state
 
 
 for msg in st.session_state.messages:
-    st.chat_message(msg["role"]).write(msg["content"])
+    with st.chat_message(msg["role"]):
+        if type(msg["content"]) is list:
+            for content in msg["content"]:
+                if content["type"] == "image_url":
+                    extension = (
+                        content["image_url"].split(";base64")[0].split("image/")[1]
+                    )
+                    base64_repr = content["image_url"].split("base64,")[1]
+                    st.image(base64.b64decode(base64_repr))
+                else:
+                    st.write(content["text"])
+        elif type(msg["content"]) is dict:
+            if msg["content"]["type"] == "image_url":
+                st.image(msg["content"]["image_url"])
+            else:
+                st.write(msg["content"]["text"])
+        elif type(msg["content"]) is str:
+            st.write(msg["content"])
+        else:
+            st.write(f"Unhandled content type: {type(msg['content'])}")
+
 
 if prompt := st.chat_input():
-    client = OpenAI(
-        base_url="http://127.0.0.1:5000/v1",
-        api_key="813",  # The OpenAI API requires an API key, but since we don't consume it, this can be any non-empty string.
-    )
+    user_message = {"role": "user", "content": [{"type": "text", "text": prompt}]}
+
+    if image_prompts:
+        for image_prompt in image_prompts:
+            extension = Path(image_prompt.name).suffix.strip(".")
+            image_bytes = image_prompt.getvalue()
+            base64_encoded = base64.b64encode(image_bytes).decode("utf-8")
+            user_message["content"].append(
+                {
+                    "type": "image_url",
+                    "image_url": f"data:image/{extension};base64,{base64_encoded}",
+                }
+            )
+    st.session_state.messages.append(user_message)
+
+    with st.chat_message("user"):
+        st.write(prompt)
+        for img in image_prompts:
+            st.image(img)
 
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    st.chat_message("user").write(prompt)
+    image_prompts = None
+    reset_per_message_state()
 
     with st.chat_message("assistant"), st.status(
         "Generating... ", expanded=True
@@ -53,15 +137,20 @@ def get_streamed_completion(completion_generator):
                 state="complete",
             )
 
-        response = st.write_stream(
-            get_streamed_completion(
-                client.chat.completions.create(
-                    model="llama3",
-                    messages=st.session_state.messages,
-                    max_tokens=response_max_tokens,
-                    stream=True,
+        try:
+            response = st.write_stream(
+                get_streamed_completion(
+                    client.chat.completions.create(
+                        model="llama3",
+                        messages=st.session_state.messages,
+                        max_tokens=response_max_tokens,
+                        temperature=temperature,
+                        stream=True,
+                    )
                 )
-            )
-        )[0]
+            )[0]
+        except Exception as e:
+            response = st.error(f"Error: {e}")
+            print(e)
 
     st.session_state.messages.append({"role": "assistant", "content": response})

torchchat/usages/openai_api.py

@@ -4,17 +4,23 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
+import base64
 import os
 import time
 import uuid
 
 from abc import ABC
 from dataclasses import dataclass
+from io import BytesIO
 from pwd import getpwuid
 from typing import Any, Dict, List, Optional, Union
 
 import torch
 
+from _torchchat_test_script import flamingo_transform, padded_collate
+from PIL import Image
+from torchtune.data import Message
+
 from torchchat.cli.download import is_model_downloaded, load_model_configs
 from torchchat.generate import Generator, GeneratorArgs
 
@@ -31,6 +37,46 @@
 # Message classes and associated objects - see the types of Messages under "Create Chat Completion >>> Request body >>> messages"
 
 
+@dataclass
+class _ContentPart(ABC):
+    """A single part of a message content field.
+
+    See the "Assistants >>> Messages >>> Create Message >>> Request body >>> content >>> Show possible types" section of the OpenAI API docs for more details.
+    """
+
+    type: str
+
+
+@dataclass
+class ImageFile:
+    file_id: str
+    detail: Optional[str]
+
+
+@dataclass
+class ImageFileContentPart(_ContentPart):
+    type: str = "image_file"
+    image_file: Optional[ImageFile] = None
+
+
+@dataclass
+class ImageUrl:
+    url: str
+    detail: Optional[str]
+
+
+@dataclass
+class ImageUrlContentPart(_ContentPart):
+    type: str = "image_url"
+    image_url: Optional[ImageUrl] = None
+
+
+@dataclass
+class TextContentPart(_ContentPart):
+    text: str = ""
+    type: str = "text"
+
+
 @dataclass
 class _AbstractMessage(ABC):
     """Base class with common parameters for message types.
@@ -42,7 +88,7 @@ class _AbstractMessage(ABC):
     """
 
     role: str
-    content: Optional[str] = None
+    content: Optional[Union[List[_ContentPart], str]] = None
 
 
 @dataclass
@@ -185,7 +231,7 @@ class ChunkDelta:
 
     tool_calls: Optional[List[ToolCall]]
     role: Optional[str]
-    content: Optional[str]
+    content: Optional[Union[List[_ContentPart], str]] = None
 
 
 @dataclass
@@ -232,18 +278,55 @@ def __init__(self, *args, **kwargs):
         """
 
         super().__init__(*args, **kwargs)
-        self.max_seq_length = (
-            self.model.config.transformer_args["text"].max_seq_length
-            + self.speculative_builder_args.speculate_k
-            + 1
-            if self.draft_model is not None
-            else self.model.config.transformer_args["text"].max_seq_length
-        )
+        self.max_seq_length = 128
+        if self.model.config.transformer_args.get("text", None):
+            self.max_seq_len = (
+                self.model.config.transformer_args["text"].max_seq_length
+                + self.speculative_builder_args.speculate_k
+                + 1
+                if self.draft_model is not None
+                else self.model.config.transformer_args["text"].max_seq_length
+            )
         # The System fingerprint is a unique identifier for the model and its configuration.
         self.system_fingerprint = (
             f"{self.builder_args.device}_{self.builder_args.precision}"
         )
 
+    def _openai_messages_to_torchtune(self, messages: List[_AbstractMessage]):
+        """Convert a list of OpenAI API messages to a list of TorchTune messages.
+
+        Args:
+            messages: A list of OpenAI API messages.
+
+        Returns:
+            A list of Torchtune Messages.
+        """
+        torchtune_messages = []
+        for message in messages:
+            torchtune_contents = []
+            if isinstance(message["content"], list):
+                for content in message["content"]:
+                    if isinstance(content, dict):
+                        if content["type"] == "image_url":
+                            torchtune_contents.append({"type": "image"})
+                        elif content["type"] == "image_file":
+                            torchtune_contents.append({"type": "image"})
+                        elif content["type"] == "text":
+                            torchtune_contents.append(
+                                {"type": "text", "content": content["text"]}
+                            )
+                    elif isinstance(content, str):
+                        torchtune_contents.append({"type": "text", "text": content})
+            else:
+                torchtune_contents.append(
+                    {"type": "text", "content": message["content"]}
+                )
+            torchtune_messages.append(
+                Message(role=message["role"], content=torchtune_contents, eot=False)
+            )
+        torchtune_messages.append(Message(role="assistant", content="", eot=False))
+        return torchtune_messages
+
     def chunked_completion(self, completion_request: CompletionRequest):
         """Handle a chat completion request and yield a chunked response.
 
@@ -271,15 +354,42 @@ def chunked_completion(self, completion_request: CompletionRequest):
         id = str(uuid.uuid4())
 
         idx = 0
-        tokens = self.chat_formatter.encode_dialog_prompt(
-            dialog=[
-                {"role": message["role"], "content": message["content"]}
-                for message in completion_request.messages
-            ]
-        )
+        images = []
 
-        encoded = torch.tensor(tokens, dtype=torch.int, device=self.builder_args.device)
-        print(self.tokenizer.decode(tokens))
+        device_sync(device=self.builder_args.device)
+        for message in completion_request.messages:
+            contents = message["content"]
+            if isinstance(contents, list):
+                for content in message["content"]:
+                    if content["type"] == "image_url":
+                        base64_decoded = base64.b64decode(
+                            content["image_url"].split(";base64,")[1]
+                        )
+                        images.append(Image.open(BytesIO(base64_decoded)))
+        print("images:", len(images), flush=True)
+        if len(images) > 0:
+            transform = flamingo_transform(str(self.tokenizer_args.tokenizer_path))
+            torchtune_messages = self._openai_messages_to_torchtune(
+                completion_request.messages
+            )
+            data = transform(
+                {"images": images, "messages": torchtune_messages}, inference=True
+            )
+            batch = padded_collate([data], self.builder_args.device)
+            batch.pop("mask")
+            encoded = batch["tokens"]
+        else:
+            tokens = self.chat_formatter.encode_dialog_prompt(
+                dialog=[
+                    {"role": message["role"], "content": message["content"]}
+                    for message in completion_request.messages
+                ]
+            )
+            print("tokens:", self.tokenizer.decode(tokens), flush=True)
+            encoded = torch.tensor(
+                tokens, dtype=torch.int, device=self.builder_args.device
+            )
+            batch = None
 
         start_pos = 0
 
@@ -293,7 +403,7 @@ def chunked_completion(self, completion_request: CompletionRequest):
             encoded_prompt=encoded,
             temperature=float(completion_request.temperature),
             chat_mode=False,
-            sequential_prefill=True,
+            sequential_prefill=False,
         )
 
         def callback(x, *, done_generating=False):
@@ -313,6 +423,7 @@ def callback(x, *, done_generating=False):
             draft_model=self.draft_model,
             speculate_k=generator_args.speculate_k,
             chat_mode=generator_args.chat_mode,
+            batch=batch,
             callback=callback,
             temperature=generator_args.temperature,
             top_k=generator_args.top_k,
@@ -323,10 +434,12 @@ def callback(x, *, done_generating=False):
         ):
             if y is None:
                 continue
+
             elif y.item() == self.tokenizer.eos_id:
                 # Stop generation if the EOS token is generated.
                 break
 
+            y = y.view(-1)
             # Decode the torch.Tensor token to a string and append to the buffer. Separate the sequences with a period token.
             content = "".join(
                 self.tokenizer.decode([self.tokenizer.encode(".")[0]] + y.tolist())[1:]