220 lines
7.6 KiB
Python
220 lines
7.6 KiB
Python
import time
|
|
from typing import Any
|
|
|
|
from mlflow.gateway.config import PaLMConfig, RouteConfig
|
|
from mlflow.gateway.exceptions import AIGatewayException
|
|
from mlflow.gateway.providers.base import BaseProvider
|
|
from mlflow.gateway.providers.utils import rename_payload_keys, send_request
|
|
from mlflow.gateway.schemas import chat, completions, embeddings
|
|
|
|
|
|
class PaLMProvider(BaseProvider):
|
|
NAME = "PaLM"
|
|
CONFIG_TYPE = PaLMConfig
|
|
|
|
def __init__(self, config: RouteConfig) -> None:
|
|
super().__init__(config)
|
|
if config.model.config is None or not isinstance(config.model.config, PaLMConfig):
|
|
raise TypeError(f"Unexpected config type {config.model.config}")
|
|
self.palm_config: PaLMConfig = config.model.config
|
|
|
|
async def _request(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
|
|
headers = {"x-goog-api-key": self.palm_config.palm_api_key}
|
|
return await send_request(
|
|
headers=headers,
|
|
base_url="https://generativelanguage.googleapis.com/v1beta3/models/",
|
|
path=path,
|
|
payload=payload,
|
|
)
|
|
|
|
async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
|
|
from fastapi.encoders import jsonable_encoder
|
|
|
|
payload = jsonable_encoder(payload, exclude_none=True)
|
|
self.check_for_model_field(payload)
|
|
if "max_tokens" in payload or "maxOutputTokens" in payload:
|
|
raise AIGatewayException(
|
|
status_code=422, detail="Max tokens is not supported for PaLM chat."
|
|
)
|
|
key_mapping = {
|
|
"stop": "stopSequences",
|
|
"n": "candidateCount",
|
|
}
|
|
for k1, k2 in key_mapping.items():
|
|
if k2 in payload:
|
|
raise AIGatewayException(
|
|
status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
|
|
)
|
|
payload = rename_payload_keys(payload, key_mapping)
|
|
# The range of PaLM's temperature is 0-1, but ours is 0-2, so we halve it
|
|
payload["temperature"] = 0.5 * payload["temperature"]
|
|
|
|
# Replace 'role' with 'author' in payload
|
|
for m in payload["messages"]:
|
|
m["author"] = m.pop("role")
|
|
|
|
# Map 'messages', 'examples, and 'context' to 'prompt'
|
|
prompt = {"messages": payload.pop("messages")}
|
|
if "examples" in payload:
|
|
prompt["examples"] = payload.pop("examples")
|
|
if "context" in payload:
|
|
prompt["context"] = payload.pop("context")
|
|
payload["prompt"] = prompt
|
|
|
|
resp = await self._request(
|
|
f"{self.config.model.name}:generateMessage",
|
|
payload,
|
|
)
|
|
|
|
# Response example
|
|
# (https://developers.generativeai.google/api/rest/generativelanguage/models/generateMessage)
|
|
# ```
|
|
# {
|
|
# "candidates": [
|
|
# {
|
|
# "author": "1",
|
|
# "content": "Hi there! How can I help you today?"
|
|
# }
|
|
# ],
|
|
# "messages": [
|
|
# {
|
|
# "author": "0",
|
|
# "content": "hi"
|
|
# }
|
|
# ]
|
|
# }
|
|
# ```
|
|
return chat.ResponsePayload(
|
|
created=int(time.time()),
|
|
model=self.config.model.name,
|
|
choices=[
|
|
chat.Choice(
|
|
index=idx,
|
|
message=chat.ResponseMessage(role=c["author"], content=c["content"]),
|
|
finish_reason=None,
|
|
)
|
|
for idx, c in enumerate(resp["candidates"])
|
|
],
|
|
usage=chat.ChatUsage(
|
|
prompt_tokens=None,
|
|
completion_tokens=None,
|
|
total_tokens=None,
|
|
),
|
|
)
|
|
|
|
async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
|
|
from fastapi.encoders import jsonable_encoder
|
|
|
|
payload = jsonable_encoder(payload, exclude_none=True)
|
|
self.check_for_model_field(payload)
|
|
key_mapping = {
|
|
"stop": "stopSequences",
|
|
"n": "candidateCount",
|
|
"max_tokens": "maxOutputTokens",
|
|
}
|
|
for k1, k2 in key_mapping.items():
|
|
if k2 in payload:
|
|
raise AIGatewayException(
|
|
status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
|
|
)
|
|
payload = rename_payload_keys(payload, key_mapping)
|
|
# The range of PaLM's temperature is 0-1, but ours is 0-2, so we halve it
|
|
payload["temperature"] = 0.5 * payload["temperature"]
|
|
payload["prompt"] = {"text": payload["prompt"]}
|
|
resp = await self._request(
|
|
f"{self.config.model.name}:generateText",
|
|
payload,
|
|
)
|
|
# Response example (https://developers.generativeai.google/api/rest/generativelanguage/models/generateText)
|
|
# ```
|
|
# {
|
|
# "candidates": [
|
|
# {
|
|
# "output": "Once upon a time, there was a young girl named Lily...",
|
|
# "safetyRatings": [
|
|
# {
|
|
# "category": "HARM_CATEGORY_DEROGATORY",
|
|
# "probability": "NEGLIGIBLE"
|
|
# }, ...
|
|
# ]
|
|
# {
|
|
# "output": "Once upon a time, there was a young boy named Billy...",
|
|
# "safetyRatings": [
|
|
# ...
|
|
# ]
|
|
# }
|
|
# ]
|
|
# }
|
|
# ```
|
|
return completions.ResponsePayload(
|
|
created=int(time.time()),
|
|
object="text_completion",
|
|
model=self.config.model.name,
|
|
choices=[
|
|
completions.Choice(
|
|
index=idx,
|
|
text=c["output"],
|
|
finish_reason=None,
|
|
)
|
|
for idx, c in enumerate(resp["candidates"])
|
|
],
|
|
usage=completions.CompletionsUsage(
|
|
prompt_tokens=None,
|
|
completion_tokens=None,
|
|
total_tokens=None,
|
|
),
|
|
)
|
|
|
|
async def embeddings(self, payload: embeddings.RequestPayload) -> embeddings.ResponsePayload:
|
|
from fastapi.encoders import jsonable_encoder
|
|
|
|
payload = jsonable_encoder(payload, exclude_none=True)
|
|
self.check_for_model_field(payload)
|
|
key_mapping = {
|
|
"input": "texts",
|
|
}
|
|
for k1, k2 in key_mapping.items():
|
|
if k2 in payload:
|
|
raise AIGatewayException(
|
|
status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
|
|
)
|
|
payload = rename_payload_keys(payload, key_mapping)
|
|
|
|
resp = await self._request(
|
|
f"{self.config.model.name}:batchEmbedText",
|
|
payload,
|
|
)
|
|
|
|
# Batch-text response example (https://developers.generativeai.google/api/rest/generativelanguage/models/batchEmbedText):
|
|
# ```
|
|
# {
|
|
# "embeddings": [
|
|
# {
|
|
# "value": [
|
|
# 3.25,
|
|
# 0.7685547,
|
|
# 2.65625,
|
|
# ...
|
|
# -0.30126953,
|
|
# -2.3554688,
|
|
# 1.2597656
|
|
# ]
|
|
# }
|
|
# ]
|
|
# }
|
|
# ```
|
|
return embeddings.ResponsePayload(
|
|
data=[
|
|
embeddings.EmbeddingObject(
|
|
embedding=embedding["value"],
|
|
index=idx,
|
|
)
|
|
for idx, embedding in enumerate(resp["embeddings"])
|
|
],
|
|
model=self.config.model.name,
|
|
usage=embeddings.EmbeddingsUsage(
|
|
prompt_tokens=None,
|
|
total_tokens=None,
|
|
),
|
|
)
|