77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
import textwrap
|
|
|
|
import gradio
|
|
|
|
from source import meta
|
|
from source.api.interface import base
|
|
from source.model.base import BaseModel
|
|
|
|
|
|
class ChatInterface(base.BaseInterface):
|
|
"""
|
|
An interface for Chat-like models.
|
|
Use the OpenAI convention (list of dict with roles and content)
|
|
"""
|
|
|
|
def __init__(self, model: "BaseModel"):
|
|
# Function to send and receive chat messages
|
|
super().__init__(model)
|
|
|
|
async def send_message(self, user_message, old_messages: list[dict], system_message: str):
|
|
# normalize the user message (the type can be wrong, especially when "edited")
|
|
if isinstance(user_message, str):
|
|
user_message: dict = {"text": user_message}
|
|
|
|
# copy the history to avoid modifying it
|
|
messages: list[dict] = old_messages.copy()
|
|
|
|
# add the system instruction
|
|
if system_message:
|
|
messages.insert(0, {"role": "system", "content": system_message})
|
|
|
|
# add the user message
|
|
# NOTE: gradio.ChatInterface add our message and the assistant message automatically
|
|
# TODO(Faraphel): add support for files - directory use user_message ? apparently, field "image" is supported.
|
|
# check "https://huggingface.co/docs/transformers/main_classes/pipelines" at "ImageTextToTextPipeline"
|
|
|
|
# TODO(Faraphel): add a "MultimodalChatInterface" to support images
|
|
messages.append({
|
|
"role": "user",
|
|
"content": user_message["text"],
|
|
})
|
|
|
|
# infer the message through the model
|
|
assistant_message = ""
|
|
async for chunk in self.model.infer(messages=messages):
|
|
assistant_message += " " + chunk.decode("utf-8")
|
|
yield assistant_message
|
|
|
|
def get_application(self):
|
|
# create a gradio interface
|
|
with gradio.Blocks(analytics_enabled=False) as application:
|
|
# header
|
|
gradio.Markdown(textwrap.dedent(f"""
|
|
# {meta.name}
|
|
## {self.model.name}
|
|
"""))
|
|
|
|
# additional settings
|
|
with gradio.Accordion("Advanced Settings") as advanced_settings:
|
|
system_prompt = gradio.Textbox(
|
|
label="System prompt",
|
|
placeholder="You are an expert in C++...",
|
|
lines=2,
|
|
)
|
|
|
|
# chat interface
|
|
gradio.ChatInterface(
|
|
fn=self.send_message,
|
|
type="messages",
|
|
multimodal=False, # TODO(Faraphel): should handle at least image and text files
|
|
editable=True,
|
|
save_history=True,
|
|
additional_inputs=[system_prompt],
|
|
additional_inputs_accordion=advanced_settings,
|
|
)
|
|
|
|
return application
|