394 lines
13 KiB
Python
394 lines
13 KiB
Python
"""
|
|
General chat utility functions.
|
|
Function 1: Basic chat completion.
|
|
Function 2: Image grounding.
|
|
"""
|
|
|
|
import re
|
|
from logging import getLogger
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Optional, Literal, Tuple
|
|
from PIL import Image, ImageDraw
|
|
from openai import OpenAIError
|
|
|
|
from .model_clients import ClientInfo, DetailedArguments
|
|
from .utils import pil_to_base64
|
|
|
|
Logger = getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ExtraParameters:
|
|
"""
|
|
Extra parameters for the chat completion.
|
|
Useful for some models with CoT.
|
|
"""
|
|
|
|
thinking: Optional[Literal["disabled", "enabled"]] = (
|
|
None # "thinking" : {"type": "disabled"} for doubao models in volcengine.
|
|
)
|
|
reasoning_effort: Optional[Literal["minimal", "low", "medium", "high"]] = (
|
|
None # "reasoning_effort" : "medium" for doubao-seed-1-6-lite-251015 or gpt-5
|
|
)
|
|
|
|
|
|
def parse_base64_image(base64_image):
|
|
"""
|
|
Parse base64 image to dict with media_type and data for anthropic.
|
|
"""
|
|
match = re.search(r"data:(image/[^;]+);base64,", base64_image)
|
|
return {
|
|
"type": "base64",
|
|
"media_type": match.group(1) if match else "image/jpeg",
|
|
"data": base64_image.split(",")[-1] if "," in base64_image else base64_image,
|
|
}
|
|
|
|
|
|
def chat_completion(
|
|
client_info: ClientInfo,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
images: list[Image.Image] | None = None,
|
|
temperature: float = 1.0,
|
|
top_p: float = 0.95,
|
|
top_k: int = 40,
|
|
max_tokens: int = 1024,
|
|
unload_after_chat: bool = True,
|
|
extra_parameters: Optional[ExtraParameters] = None,
|
|
) -> str:
|
|
"""
|
|
Basic chat completion.
|
|
"""
|
|
# Basic check of the user's input.
|
|
if images is None:
|
|
payload_images: list = []
|
|
elif len(images) > 4:
|
|
Logger.warning(
|
|
"Warning: The number of images is greater than 4. Only the first 4 images will be used."
|
|
)
|
|
payload_images: list = images[:4]
|
|
else:
|
|
payload_images: list = images
|
|
|
|
base64_images: list[str] = [pil_to_base64(image) for image in payload_images]
|
|
orig_args: DetailedArguments = client_info.arguments
|
|
messages: list[dict] = [
|
|
{
|
|
"role": "system",
|
|
"content": system_prompt,
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": user_prompt},
|
|
*[
|
|
{"type": "image_url", "image_url": {"url": base64_image}}
|
|
for base64_image in base64_images
|
|
],
|
|
],
|
|
},
|
|
]
|
|
# prepare the arguments
|
|
# Due to ollama issue (Cannot unload model using openai chat api.).
|
|
# We have to generate a single payload here.
|
|
if client_info.client_type == "ollama":
|
|
payload = {
|
|
"model": orig_args.model,
|
|
"messages": [
|
|
{
|
|
"role": "system",
|
|
"content": system_prompt,
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": user_prompt,
|
|
"images": [base64_image for base64_image in base64_images],
|
|
},
|
|
],
|
|
"options": {
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"top_k": top_k,
|
|
"num_predict": max_tokens,
|
|
},
|
|
}
|
|
if extra_parameters:
|
|
# "minimal" thinking effect is not supported by ollama.
|
|
# but "minimal" thinking effect belong to no think is volcengine.
|
|
need_think = extra_parameters.thinking == "enabled"
|
|
payload["think"] = need_think
|
|
if need_think and extra_parameters.reasoning_effort != "minimal":
|
|
payload["think"] = extra_parameters.reasoning_effort
|
|
if unload_after_chat:
|
|
payload["keep_alive"] = "0"
|
|
|
|
elif client_info.client_type == "openai":
|
|
payload = {
|
|
"model": orig_args.model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
if extra_parameters:
|
|
if extra_parameters.thinking == "enabled":
|
|
if "seed-1-6" in orig_args.model:
|
|
payload["extra_body"] = {}
|
|
payload["extra_body"]["thinking"] = {
|
|
"type": extra_parameters.thinking
|
|
}
|
|
payload["reasoning_effort"] = extra_parameters.reasoning_effort
|
|
|
|
elif client_info.client_type == "openai-responses":
|
|
payload = {
|
|
"model": orig_args.model,
|
|
"input": [
|
|
{
|
|
"role": "system",
|
|
"content": system_prompt,
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "input_text", "text": user_prompt},
|
|
*[
|
|
{"type": "input_image", "image_url": base64_image}
|
|
for base64_image in base64_images
|
|
],
|
|
],
|
|
},
|
|
],
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"max_output_tokens": max_tokens,
|
|
}
|
|
if extra_parameters:
|
|
model: str = orig_args.model
|
|
if "seed-1-6" in model:
|
|
payload["extra_body"] = {}
|
|
payload["extra_body"]["thinking"] = {"type": extra_parameters.thinking}
|
|
else:
|
|
# volcengine does not support thinking effort in responses api
|
|
payload["reasoning"] = {}
|
|
payload["reasoning"]["effort"] = extra_parameters.reasoning_effort
|
|
|
|
elif client_info.client_type == "mistral":
|
|
payload = {
|
|
"model": orig_args.model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"top_k": top_k,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
|
|
elif client_info.client_type == "anthropic":
|
|
anthropic_messages = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": user_prompt},
|
|
*[
|
|
{
|
|
"type": "image",
|
|
"source": parse_base64_image(base64_image),
|
|
}
|
|
for base64_image in base64_images
|
|
],
|
|
],
|
|
},
|
|
]
|
|
payload = {
|
|
"model": orig_args.model,
|
|
"system": [{"text": system_prompt, "type": "text"}],
|
|
"messages": anthropic_messages,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"top_k": top_k,
|
|
"max_tokens": max_tokens,
|
|
"thinking": {"type": "disabled"},
|
|
}
|
|
if extra_parameters and extra_parameters.thinking == "enabled":
|
|
payload["thinking"] = {"type": extra_parameters.thinking}
|
|
if (
|
|
extra_parameters.reasoning_effort
|
|
and extra_parameters.reasoning_effort != "minimal"
|
|
):
|
|
effort_token = {
|
|
"low": 1024,
|
|
"medium": 2048,
|
|
"high": 4096,
|
|
}
|
|
payload["thinking"]["budget_tokens"] = effort_token.get( # type: ignore
|
|
extra_parameters.reasoning_effort, 1024
|
|
)
|
|
else:
|
|
payload["thinking"] = {"type": "disabled"}
|
|
|
|
else:
|
|
raise ValueError("The client type is not supported.")
|
|
|
|
# Call the API
|
|
try:
|
|
response = client_info.chat_func(**payload)
|
|
except OpenAIError as e:
|
|
raise RuntimeError(f"Error in openai chat completion: {e}") from e
|
|
except Exception as e:
|
|
raise RuntimeError(f"Unexpected error in chat completion: {e}") from e
|
|
|
|
# Parse the response
|
|
if client_info.client_type in ["openai", "mistral"]: # Chat Completion API
|
|
result: str = response.choices[0].message.content
|
|
elif client_info.client_type == "openai-responses": # Responses API
|
|
result: str = response.output_text
|
|
elif client_info.client_type == "ollama":
|
|
result: str = response.message.content
|
|
elif client_info.client_type == "anthropic":
|
|
result: str = response.content[0].text
|
|
else:
|
|
raise ValueError("The message type is not supported.")
|
|
return result
|
|
|
|
|
|
def grounding(
|
|
client_info: ClientInfo,
|
|
item: Optional[str],
|
|
image: Image.Image,
|
|
thinking: Literal["disabled", "enabled"] = "disabled",
|
|
mode: Literal["minimal", "low", "medium", "high"] = "minimal",
|
|
unload_after_chat: bool = True,
|
|
) -> Tuple[list[list[int]], Image.Image]:
|
|
"""
|
|
Image grounding with item.
|
|
"""
|
|
# Basic check of the user's input.
|
|
if item is None:
|
|
raise ValueError("The item is None")
|
|
|
|
# prepare the prompt
|
|
system_prompt = "You are a professional image grounding assistant."
|
|
user_prompt = (
|
|
f"Please locate the item '{item}' in the image accurately. "
|
|
+ "Response in coordinate of the bounding box. "
|
|
+ "The format is <bbox>x_min y_min x_max y_max</bbox> in percentage(0-1000). "
|
|
+ "If there are multiple items, please list all bounding boxes. "
|
|
)
|
|
|
|
# prepare the extra parameters
|
|
extra_parameters = ExtraParameters(thinking=thinking, reasoning_effort=mode)
|
|
|
|
# call the API
|
|
response: str = chat_completion(
|
|
client_info=client_info,
|
|
system_prompt=system_prompt,
|
|
user_prompt=user_prompt,
|
|
images=[image],
|
|
temperature=1.0,
|
|
top_p=0.7,
|
|
max_tokens=4096,
|
|
unload_after_chat=unload_after_chat,
|
|
extra_parameters=extra_parameters,
|
|
)
|
|
|
|
# responses: <bbox>0 55 63 250</bbox><bbox>0 300 51 427</bbox>...
|
|
bboxes: list[list[int]] = []
|
|
for bbox in re.findall(r"<bbox>(.*?)</bbox>", response):
|
|
# seed 1.6
|
|
bbox = [int(x) for x in bbox.split()]
|
|
if len(bbox) != 4:
|
|
raise ValueError("BBox is invaild. Please retry again.")
|
|
x_min, y_min, x_max, y_max = tuple(bbox)
|
|
w, h = image.size
|
|
x_min_real = int(x_min * w / 1000)
|
|
y_min_real = int(y_min * h / 1000)
|
|
x_max_real = int(x_max * w / 1000)
|
|
y_max_real = int(y_max * h / 1000)
|
|
|
|
bbox_x = x_min_real
|
|
bbox_y = y_min_real
|
|
bbox_w = x_max_real - x_min_real
|
|
bbox_h = y_max_real - y_min_real
|
|
|
|
bboxes.append([bbox_x, bbox_y, bbox_w, bbox_h])
|
|
|
|
# Create a copy of the image to draw on
|
|
draw_image = image.copy()
|
|
draw = ImageDraw.Draw(draw_image)
|
|
|
|
# Define bbox style
|
|
bbox_color = "red" # Red color for bboxes
|
|
bbox_width = 2 # Line width
|
|
|
|
# Draw each bbox
|
|
for bbox in bboxes:
|
|
if len(bbox) == 4: # Ensure bbox has 4 coordinates [x1, y1, x2, y2]
|
|
x1, y1, x2, y2 = bbox
|
|
# Draw rectangle
|
|
# Pillow uses [x1, x2, x3, x4]
|
|
draw.rectangle(
|
|
[x1, y1, x1 + x2, y1 + y2], outline=bbox_color, width=bbox_width
|
|
)
|
|
|
|
return bboxes, draw_image
|
|
|
|
|
|
def captioning(
|
|
client_info: ClientInfo,
|
|
image: Image.Image,
|
|
language: str = "English",
|
|
num_max_sentences: int = 10,
|
|
thinking: Literal["disabled", "enabled"] = "disabled",
|
|
mode: Literal["minimal", "low", "medium", "high"] = "minimal",
|
|
unload_after_chat: bool = True,
|
|
) -> str:
|
|
"""
|
|
Image captioning.
|
|
"""
|
|
|
|
# Captioning prompt based on Qwen-Image arxiv paper.
|
|
# https://arxiv.org/pdf/2508.02324
|
|
# Figure 12 in Page 13.
|
|
# Modified by Deepseek-V3.2-Exp to remove unused json format.
|
|
# Better for using thinking mode with "enabled" and "medium" reasoning_effort.
|
|
system_prompt = (
|
|
"As a professional image annotator, "
|
|
+ "your primary task is to generate a detailed and natural caption for the input image, "
|
|
+ "focusing on authenticity and accuracy without any generalizations. "
|
|
+ "Write the caption in descriptive, flowing text, "
|
|
+ "avoiding structured formats or rich text elements. "
|
|
+ "Enrich the description by including specific object attributes, "
|
|
+ "visual relationships between objects, and environmental context "
|
|
+ "to provide a comprehensive view. "
|
|
+ "If any text is visible in the image, identify it exactly as seen "
|
|
+ "and highlight it within the caption using quotation marks, "
|
|
+ "without translating or explaining the content. "
|
|
+ "Ensure all details are grounded in what is actually present, "
|
|
+ "maintaining a truthful representation of the scene. "
|
|
+ "This approach guarantees a clear and informative caption "
|
|
+ "that captures the essence of the image effectively."
|
|
)
|
|
user_prompt = (
|
|
"Here's the user's input. "
|
|
+ f"Please generate a detailed caption for the image in language {language} "
|
|
+ f"with maxium {num_max_sentences} sentences."
|
|
)
|
|
|
|
# prepare the extra parameters
|
|
extra_parameters = ExtraParameters(thinking=thinking, reasoning_effort=mode)
|
|
|
|
# call the API
|
|
response: str = chat_completion(
|
|
client_info=client_info,
|
|
system_prompt=system_prompt,
|
|
user_prompt=user_prompt,
|
|
images=[image],
|
|
temperature=1.0,
|
|
top_p=0.7,
|
|
max_tokens=4096,
|
|
unload_after_chat=unload_after_chat,
|
|
extra_parameters=extra_parameters,
|
|
)
|
|
|
|
return response
|