mirror of https://github.com/HazyResearch/manifest
Laurel Orr
1 year ago
committed by
GitHub
parent
4903c7e7e8
commit
b52a4d9a4b
@ -0,0 +1,149 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%load_ext autoreload\n",
|
||||
"%autoreload 2"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"AZURE_KEY = \"API_KEY::URL\"\n",
|
||||
"OPENAI_KEY = \"sk-XXX\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Use Azure and OpenAI models"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from manifest import Manifest\n",
|
||||
"from manifest.connections.client_pool import ClientConnection\n",
|
||||
"from pathlib import Path\n",
|
||||
"\n",
|
||||
"cache_path = Path(\"manifest.db\")\n",
|
||||
"if cache_path.exists():\n",
|
||||
" cache_path.unlink()\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"azure = ClientConnection(\n",
|
||||
" client_name=\"azureopenai\",\n",
|
||||
" client_connection=AZURE_KEY,\n",
|
||||
" engine=\"text-davinci-003\",\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"manifest = Manifest(client_pool=[azure], \n",
|
||||
" cache_name=\"sqlite\",\n",
|
||||
" cache_connection=\"manifest.db\"\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"openai = ClientConnection(\n",
|
||||
" client_name=\"openai\",\n",
|
||||
" client_connection=OPENAI_KEY,\n",
|
||||
" engine=\"text-davinci-003\",\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"manifest_openai_nocache = Manifest(client_pool=[openai])\n",
|
||||
"\n",
|
||||
"manifest_openai = Manifest(client_pool=[openai], \n",
|
||||
" cache_name=\"sqlite\",\n",
|
||||
" cache_connection=\"manifest.db\"\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Show caches are the same\n",
|
||||
"text = \"What is the meaning of life?\"\n",
|
||||
"res = manifest.run(text, max_tokens=100, temperature=0.7, return_response=True)\n",
|
||||
"print(res.get_response())\n",
|
||||
"print(res.is_cached())\n",
|
||||
"res2 = manifest_openai.run(text, max_tokens=100, temperature=0.7, return_response=True)\n",
|
||||
"print(res2.is_cached())\n",
|
||||
"\n",
|
||||
"assert res2.get_response() == res.get_response()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"azure_chat = ClientConnection(\n",
|
||||
" client_name=\"azureopenaichat\",\n",
|
||||
" client_connection=AZURE_KEY,\n",
|
||||
" engine=\"gpt-3.5-turbo\",\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"manifest = Manifest(client_pool=[azure_chat])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"print(manifest.run(\"What do you think is the best food?\", max_tokens=100))\n",
|
||||
"\n",
|
||||
"chat_dict = [\n",
|
||||
" {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
|
||||
" {\"role\": \"user\", \"content\": \"Who won the world series in 2020?\"},\n",
|
||||
" {\"role\": \"assistant\", \"content\": \"The Los Angeles Dodgers won the World Series in 2020.\"},\n",
|
||||
" {\"role\": \"user\", \"content\": \"Where was it played?\"}\n",
|
||||
"]\n",
|
||||
"print(manifest.run(chat_dict, max_tokens=100))"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "manifest",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.4"
|
||||
},
|
||||
"orig_nbformat": 4,
|
||||
"vscode": {
|
||||
"interpreter": {
|
||||
"hash": "fddffe4ac3b9f00470127629076101c1b5f38ecb1e7358b567d19305425e9491"
|
||||
}
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
@ -0,0 +1,117 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%load_ext autoreload\n",
|
||||
"%autoreload 2"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"GOOGLE_KEY = \"KEY::PROJECT_ID\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Use GoogleVertexAPI"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from manifest import Manifest\n",
|
||||
"from manifest.connections.client_pool import ClientConnection\n",
|
||||
"\n",
|
||||
"google_bison = ClientConnection(\n",
|
||||
" client_name=\"google\",\n",
|
||||
" client_connection=GOOGLE_KEY\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"manifest = Manifest(client_pool=[google_bison])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Simple question\n",
|
||||
"print(manifest.run(\"What is your name\", max_tokens=40))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from manifest import Manifest\n",
|
||||
"from manifest.connections.client_pool import ClientConnection\n",
|
||||
"\n",
|
||||
"google_bison = ClientConnection(\n",
|
||||
" client_name=\"googlechat\",\n",
|
||||
" client_connection=GOOGLE_KEY\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"manifest = Manifest(client_pool=[google_bison])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"chat_dict = [\n",
|
||||
" # {\"author\": \"bot\", \"content\": \"You are a helpful assistant.\"},\n",
|
||||
" {\"author\": \"user\", \"content\": \"Who won the world series in 2020?\"},\n",
|
||||
" {\"author\": \"bot\", \"content\": \"The Los Angeles Dodgers won the World Series in 2020.\"},\n",
|
||||
" {\"author\": \"user\", \"content\": \"Where was it played?\"}\n",
|
||||
"]\n",
|
||||
"print(manifest.run(chat_dict, max_tokens=8))"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "manifest",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.4"
|
||||
},
|
||||
"orig_nbformat": 4,
|
||||
"vscode": {
|
||||
"interpreter": {
|
||||
"hash": "fddffe4ac3b9f00470127629076101c1b5f38ecb1e7358b567d19305425e9491"
|
||||
}
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
@ -0,0 +1,111 @@
|
||||
"""OpenAI client."""
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, Optional, Type
|
||||
|
||||
from manifest.clients.openai import OPENAI_ENGINES, OpenAIClient
|
||||
from manifest.request import LMRequest, Request
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Azure deployment name can only use letters and numbers, no spaces. Hyphens ("-") and
|
||||
# underscores ("_") may be used, except as ending characters. We create this mapping to
|
||||
# handle difference between Azure and OpenAI
|
||||
AZURE_DEPLOYMENT_NAME_MAPPING = {
|
||||
"gpt-3.5-turbo": "gpt-35-turbo",
|
||||
"gpt-3.5-turbo-0301": "gpt-35-turbo-0301",
|
||||
}
|
||||
OPENAI_DEPLOYMENT_NAME_MAPPING = {
|
||||
"gpt-35-turbo": "gpt-3.5-turbo",
|
||||
"gpt-35-turbo-0301": "gpt-3.5-turbo-0301",
|
||||
}
|
||||
|
||||
|
||||
class AzureClient(OpenAIClient):
|
||||
"""Azure client."""
|
||||
|
||||
PARAMS = OpenAIClient.PARAMS
|
||||
REQUEST_CLS: Type[Request] = LMRequest
|
||||
NAME = "azureopenai"
|
||||
|
||||
def connect(
|
||||
self,
|
||||
connection_str: Optional[str] = None,
|
||||
client_args: Dict[str, Any] = {},
|
||||
) -> None:
|
||||
"""
|
||||
Connect to the AzureOpenAI server.
|
||||
|
||||
connection_str is passed as default AZURE_OPENAI_KEY if variable not set.
|
||||
|
||||
Args:
|
||||
connection_str: connection string.
|
||||
client_args: client arguments.
|
||||
"""
|
||||
connection_parts = connection_str.split("::")
|
||||
if len(connection_parts) == 1:
|
||||
self.api_key = connection_parts[0]
|
||||
elif len(connection_parts) == 2:
|
||||
self.api_key, self.host = connection_parts
|
||||
else:
|
||||
raise ValueError(
|
||||
"Invalid connection string. "
|
||||
"Must be either AZURE_OPENAI_KEY or "
|
||||
"AZURE_OPENAI_KEY::AZURE_OPENAI_ENDPOINT"
|
||||
)
|
||||
self.api_key = self.api_key or os.environ.get("AZURE_OPENAI_KEY")
|
||||
if self.api_key is None:
|
||||
raise ValueError(
|
||||
"AzureOpenAI API key not set. Set AZURE_OPENAI_KEY environment "
|
||||
"variable or pass through `client_connection`."
|
||||
)
|
||||
self.host = self.host or os.environ.get("AZURE_OPENAI_ENDPOINT")
|
||||
self.host = self.host.rstrip("/")
|
||||
if self.host is None:
|
||||
raise ValueError(
|
||||
"Azure Service URL not set "
|
||||
"(e.g. https://openai-azure-service.openai.azure.com/)."
|
||||
" Set AZURE_OPENAI_ENDPOINT or pass through `client_connection`."
|
||||
" as AZURE_OPENAI_KEY::AZURE_OPENAI_ENDPOINT"
|
||||
)
|
||||
for key in self.PARAMS:
|
||||
setattr(self, key, client_args.pop(key, self.PARAMS[key][1]))
|
||||
if getattr(self, "engine") not in OPENAI_ENGINES:
|
||||
raise ValueError(
|
||||
f"Invalid engine {getattr(self, 'engine')}. Must be {OPENAI_ENGINES}."
|
||||
)
|
||||
|
||||
def get_generation_url(self) -> str:
|
||||
"""Get generation URL."""
|
||||
engine = getattr(self, "engine")
|
||||
deployment_name = AZURE_DEPLOYMENT_NAME_MAPPING.get(engine, engine)
|
||||
return (
|
||||
self.host
|
||||
+ "/openai/deployments/"
|
||||
+ deployment_name
|
||||
+ "/completions?api-version=2023-05-15"
|
||||
)
|
||||
|
||||
def get_generation_header(self) -> Dict[str, str]:
|
||||
"""
|
||||
Get generation header.
|
||||
|
||||
Returns:
|
||||
header.
|
||||
"""
|
||||
return {"api-key": f"{self.api_key}"}
|
||||
|
||||
def get_model_params(self) -> Dict:
|
||||
"""
|
||||
Get model params.
|
||||
|
||||
By getting model params from the server, we can add to request
|
||||
and make sure cache keys are unique to model.
|
||||
|
||||
Returns:
|
||||
model params.
|
||||
"""
|
||||
# IMPORTANT!!!
|
||||
# Azure models are the same as openai models. So we want to unify their
|
||||
# cached. Make sure we retrun the OpenAI name here.
|
||||
return {"model_name": OpenAIClient.NAME, "engine": getattr(self, "engine")}
|
||||
@ -0,0 +1,114 @@
|
||||
"""OpenAI client."""
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from manifest.clients.openai_chat import OPENAICHAT_ENGINES, OpenAIChatClient
|
||||
from manifest.request import LMRequest
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Azure deployment name can only use letters and numbers, no spaces. Hyphens ("-") and
|
||||
# underscores ("_") may be used, except as ending characters. We create this mapping to
|
||||
# handle difference between Azure and OpenAI
|
||||
AZURE_DEPLOYMENT_NAME_MAPPING = {
|
||||
"gpt-3.5-turbo": "gpt-35-turbo",
|
||||
"gpt-3.5-turbo-0301": "gpt-35-turbo-0301",
|
||||
}
|
||||
OPENAI_DEPLOYMENT_NAME_MAPPING = {
|
||||
"gpt-35-turbo": "gpt-3.5-turbo",
|
||||
"gpt-35-turbo-0301": "gpt-3.5-turbo-0301",
|
||||
}
|
||||
|
||||
|
||||
class AzureChatClient(OpenAIChatClient):
|
||||
"""Azure chat client."""
|
||||
|
||||
# User param -> (client param, default value)
|
||||
PARAMS = OpenAIChatClient.PARAMS
|
||||
REQUEST_CLS = LMRequest
|
||||
NAME = "azureopenaichat"
|
||||
IS_CHAT = True
|
||||
|
||||
def connect(
|
||||
self,
|
||||
connection_str: Optional[str] = None,
|
||||
client_args: Dict[str, Any] = {},
|
||||
) -> None:
|
||||
"""
|
||||
Connect to the AzureOpenAI server.
|
||||
|
||||
connection_str is passed as default AZURE_OPENAI_KEY if variable not set.
|
||||
|
||||
Args:
|
||||
connection_str: connection string.
|
||||
client_args: client arguments.
|
||||
"""
|
||||
connection_parts = connection_str.split("::")
|
||||
if len(connection_parts) == 1:
|
||||
self.api_key = connection_parts[0]
|
||||
elif len(connection_parts) == 2:
|
||||
self.api_key, self.host = connection_parts
|
||||
else:
|
||||
raise ValueError(
|
||||
"Invalid connection string. "
|
||||
"Must be either AZURE_OPENAI_KEY or "
|
||||
"AZURE_OPENAI_KEY::AZURE_OPENAI_ENDPOINT"
|
||||
)
|
||||
self.api_key = self.api_key or os.environ.get("AZURE_OPENAI_KEY")
|
||||
if self.api_key is None:
|
||||
raise ValueError(
|
||||
"AzureOpenAI API key not set. Set AZURE_OPENAI_KEY environment "
|
||||
"variable or pass through `client_connection`."
|
||||
)
|
||||
self.host = self.host or os.environ.get("AZURE_OPENAI_ENDPOINT")
|
||||
self.host = self.host.rstrip("/")
|
||||
if self.host is None:
|
||||
raise ValueError(
|
||||
"Azure Service URL not set "
|
||||
"(e.g. https://openai-azure-service.openai.azure.com/)."
|
||||
" Set AZURE_OPENAI_ENDPOINT or pass through `client_connection`."
|
||||
" as AZURE_OPENAI_KEY::AZURE_OPENAI_ENDPOINT"
|
||||
)
|
||||
for key in self.PARAMS:
|
||||
setattr(self, key, client_args.pop(key, self.PARAMS[key][1]))
|
||||
if getattr(self, "engine") not in OPENAICHAT_ENGINES:
|
||||
raise ValueError(
|
||||
f"Invalid engine {getattr(self, 'engine')}. "
|
||||
f"Must be {OPENAICHAT_ENGINES}."
|
||||
)
|
||||
|
||||
def get_generation_url(self) -> str:
|
||||
"""Get generation URL."""
|
||||
engine = getattr(self, "engine")
|
||||
deployment_name = AZURE_DEPLOYMENT_NAME_MAPPING.get(engine, engine)
|
||||
return (
|
||||
self.host
|
||||
+ "/openai/deployments/"
|
||||
+ deployment_name
|
||||
+ "/chat/completions?api-version=2023-05-15"
|
||||
)
|
||||
|
||||
def get_generation_header(self) -> Dict[str, str]:
|
||||
"""
|
||||
Get generation header.
|
||||
|
||||
Returns:
|
||||
header.
|
||||
"""
|
||||
return {"api-key": f"{self.api_key}"}
|
||||
|
||||
def get_model_params(self) -> Dict:
|
||||
"""
|
||||
Get model params.
|
||||
|
||||
By getting model params from the server, we can add to request
|
||||
and make sure cache keys are unique to model.
|
||||
|
||||
Returns:
|
||||
model params.
|
||||
"""
|
||||
# IMPORTANT!!!
|
||||
# Azure models are the same as openai models. So we want to unify their
|
||||
# cached. Make sure we retrun the OpenAI name here.
|
||||
return {"model_name": OpenAIChatClient.NAME, "engine": getattr(self, "engine")}
|
||||
@ -0,0 +1,190 @@
|
||||
"""OpenAI client."""
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
from typing import Any, Dict, Optional, Type
|
||||
|
||||
from manifest.clients.client import Client
|
||||
from manifest.request import LMRequest, Request
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# https://cloud.google.com/vertex-ai/docs/generative-ai/start/quickstarts/api-quickstart
|
||||
GOOGLE_ENGINES = {
|
||||
"text-bison",
|
||||
}
|
||||
|
||||
|
||||
def get_project_id() -> Optional[str]:
|
||||
"""Get project ID.
|
||||
|
||||
Run
|
||||
`gcloud config get-value project`
|
||||
"""
|
||||
try:
|
||||
project_id = subprocess.run(
|
||||
["gcloud", "config", "get-value", "project"],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
if project_id.stderr.decode("utf-8").strip():
|
||||
return None
|
||||
return project_id.stdout.decode("utf-8").strip()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
class GoogleClient(Client):
|
||||
"""Google client."""
|
||||
|
||||
# User param -> (client param, default value)
|
||||
PARAMS = {
|
||||
"engine": ("model", "text-bison"),
|
||||
"temperature": ("temperature", 1.0),
|
||||
"max_tokens": ("maxOutputTokens", 10),
|
||||
"top_p": ("topP", 1.0),
|
||||
"top_k": ("topK", 1),
|
||||
"batch_size": ("batch_size", 20),
|
||||
}
|
||||
REQUEST_CLS: Type[Request] = LMRequest
|
||||
NAME = "google"
|
||||
|
||||
def connect(
|
||||
self,
|
||||
connection_str: Optional[str] = None,
|
||||
client_args: Dict[str, Any] = {},
|
||||
) -> None:
|
||||
"""
|
||||
Connect to the GoogleVertex API.
|
||||
|
||||
connection_str is passed as default GOOGLE_API_KEY if variable not set.
|
||||
|
||||
Args:
|
||||
connection_str: connection string.
|
||||
client_args: client arguments.
|
||||
"""
|
||||
connection_parts = connection_str.split("::")
|
||||
if len(connection_parts) == 1:
|
||||
self.api_key = connection_parts[0]
|
||||
self.project_id = None
|
||||
elif len(connection_parts) == 2:
|
||||
self.api_key, self.project_id = connection_parts
|
||||
else:
|
||||
raise ValueError(
|
||||
"Invalid connection string. "
|
||||
"Must be either API_KEY or API_KEY::PROJECT_ID"
|
||||
)
|
||||
self.api_key = self.api_key or os.environ.get("GOOGLE_API_KEY")
|
||||
if self.api_key is None:
|
||||
raise ValueError(
|
||||
"GoogleVertex API key not set. Set GOOGLE_API_KEY environment "
|
||||
"variable or pass through `client_connection`. This can be "
|
||||
"found by running `gcloud auth print-access-token`"
|
||||
)
|
||||
self.project_id = (
|
||||
self.project_id or os.environ.get("GOOGLE_PROJECT_ID") or get_project_id()
|
||||
)
|
||||
if self.project_id is None:
|
||||
raise ValueError("GoogleVertex project ID not set. Set GOOGLE_PROJECT_ID")
|
||||
self.host = f"https://us-central1-aiplatform.googleapis.com/v1/projects/{self.project_id}/locations/us-central1/publishers/google/models" # noqa: E501
|
||||
|
||||
for key in self.PARAMS:
|
||||
setattr(self, key, client_args.pop(key, self.PARAMS[key][1]))
|
||||
if getattr(self, "engine") not in GOOGLE_ENGINES:
|
||||
raise ValueError(
|
||||
f"Invalid engine {getattr(self, 'engine')}. Must be {GOOGLE_ENGINES}."
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the client."""
|
||||
pass
|
||||
|
||||
def get_generation_url(self) -> str:
|
||||
"""Get generation URL."""
|
||||
model = getattr(self, "engine")
|
||||
return self.host + f"/{model}:predict"
|
||||
|
||||
def get_generation_header(self) -> Dict[str, str]:
|
||||
"""
|
||||
Get generation header.
|
||||
|
||||
Returns:
|
||||
header.
|
||||
"""
|
||||
return {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
def supports_batch_inference(self) -> bool:
|
||||
"""Return whether the client supports batch inference."""
|
||||
return True
|
||||
|
||||
def get_model_params(self) -> Dict:
|
||||
"""
|
||||
Get model params.
|
||||
|
||||
By getting model params from the server, we can add to request
|
||||
and make sure cache keys are unique to model.
|
||||
|
||||
Returns:
|
||||
model params.
|
||||
"""
|
||||
return {"model_name": self.NAME, "engine": getattr(self, "engine")}
|
||||
|
||||
def preprocess_request_params(self, request: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Preprocess request params.
|
||||
|
||||
Args:
|
||||
request: request params.
|
||||
|
||||
Returns:
|
||||
request params.
|
||||
"""
|
||||
# Refortmat the request params for google
|
||||
prompt = request.pop("prompt")
|
||||
if isinstance(prompt, str):
|
||||
prompt_list = [prompt]
|
||||
else:
|
||||
prompt_list = prompt
|
||||
google_request = {
|
||||
"instances": [{"prompt": prompt} for prompt in prompt_list],
|
||||
"parameters": request,
|
||||
}
|
||||
return super().preprocess_request_params(google_request)
|
||||
|
||||
def postprocess_response(
|
||||
self, response: Dict[str, Any], request: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Validate response as dict.
|
||||
|
||||
Assumes response is dict
|
||||
{
|
||||
"predictions": [
|
||||
{
|
||||
"safetyAttributes": {
|
||||
"categories": ["Violent", "Sexual"],
|
||||
"blocked": false,
|
||||
"scores": [0.1, 0.1]
|
||||
},
|
||||
"content": "SELECT * FROM "WWW";"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
Args:
|
||||
response: response
|
||||
request: request
|
||||
|
||||
Return:
|
||||
response as dict
|
||||
"""
|
||||
google_predictions = response.pop("predictions")
|
||||
new_response = {
|
||||
"choices": [
|
||||
{
|
||||
"text": prediction["content"],
|
||||
}
|
||||
for prediction in google_predictions
|
||||
]
|
||||
}
|
||||
return super().postprocess_response(new_response, request)
|
||||
@ -0,0 +1,155 @@
|
||||
"""OpenAI client."""
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, Optional, Type
|
||||
|
||||
from manifest.clients.google import GoogleClient, get_project_id
|
||||
from manifest.request import LMRequest, Request
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# https://cloud.google.com/vertex-ai/docs/generative-ai/start/quickstarts/api-quickstart
|
||||
GOOGLE_ENGINES = {
|
||||
"chat-bison",
|
||||
}
|
||||
|
||||
|
||||
class GoogleChatClient(GoogleClient):
|
||||
"""GoogleChat client."""
|
||||
|
||||
# User param -> (client param, default value)
|
||||
PARAMS = {
|
||||
"engine": ("model", "chat-bison"),
|
||||
"temperature": ("temperature", 1.0),
|
||||
"max_tokens": ("maxOutputTokens", 10),
|
||||
"top_p": ("topP", 1.0),
|
||||
"top_k": ("topK", 1),
|
||||
"batch_size": ("batch_size", 20),
|
||||
}
|
||||
REQUEST_CLS: Type[Request] = LMRequest
|
||||
NAME = "googlechat"
|
||||
IS_CHAT = True
|
||||
|
||||
def connect(
|
||||
self,
|
||||
connection_str: Optional[str] = None,
|
||||
client_args: Dict[str, Any] = {},
|
||||
) -> None:
|
||||
"""
|
||||
Connect to the GoogleVertex API.
|
||||
|
||||
connection_str is passed as default GOOGLE_API_KEY if variable not set.
|
||||
|
||||
Args:
|
||||
connection_str: connection string.
|
||||
client_args: client arguments.
|
||||
"""
|
||||
connection_parts = connection_str.split("::")
|
||||
if len(connection_parts) == 1:
|
||||
self.api_key = connection_parts[0]
|
||||
elif len(connection_parts) == 2:
|
||||
self.api_key, self.project_id = connection_parts
|
||||
else:
|
||||
raise ValueError(
|
||||
"Invalid connection string. "
|
||||
"Must be either API_KEY or API_KEY::PROJECT_ID"
|
||||
)
|
||||
self.api_key = self.api_key or os.environ.get("GOOGLE_API_KEY")
|
||||
if self.api_key is None:
|
||||
raise ValueError(
|
||||
"GoogleVertex API key not set. Set GOOGLE_API_KEY environment "
|
||||
"variable or pass through `client_connection`. This can be "
|
||||
"found by running `gcloud auth print-access-token`"
|
||||
)
|
||||
self.project_id = (
|
||||
self.project_id or os.environ.get("GOOGLE_PROJECT_ID") or get_project_id()
|
||||
)
|
||||
if self.project_id is None:
|
||||
raise ValueError("GoogleVertex project ID not set. Set GOOGLE_PROJECT_ID")
|
||||
self.host = f"https://us-central1-aiplatform.googleapis.com/v1/projects/{self.project_id}/locations/us-central1/publishers/google/models" # noqa: E501
|
||||
|
||||
for key in self.PARAMS:
|
||||
setattr(self, key, client_args.pop(key, self.PARAMS[key][1]))
|
||||
if getattr(self, "engine") not in GOOGLE_ENGINES:
|
||||
raise ValueError(
|
||||
f"Invalid engine {getattr(self, 'engine')}. Must be {GOOGLE_ENGINES}."
|
||||
)
|
||||
|
||||
def supports_batch_inference(self) -> bool:
|
||||
"""Return whether the client supports batch inference."""
|
||||
return False
|
||||
|
||||
def preprocess_request_params(self, request: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Preprocess request params.
|
||||
|
||||
Args:
|
||||
request: request params.
|
||||
|
||||
Returns:
|
||||
request params.
|
||||
"""
|
||||
# Format for chat model
|
||||
request = copy.deepcopy(request)
|
||||
prompt = request.pop("prompt")
|
||||
if isinstance(prompt, str):
|
||||
messages = [{"author": "user", "content": prompt}]
|
||||
elif isinstance(prompt, list) and isinstance(prompt[0], str):
|
||||
prompt_list = prompt
|
||||
messages = [{"author": "user", "content": prompt} for prompt in prompt_list]
|
||||
elif isinstance(prompt, list) and isinstance(prompt[0], dict):
|
||||
for pmt_dict in prompt:
|
||||
if "author" not in pmt_dict or "content" not in pmt_dict:
|
||||
raise ValueError(
|
||||
"Prompt must be list of dicts with 'author' and 'content' "
|
||||
f"keys. Got {prompt}."
|
||||
)
|
||||
messages = prompt
|
||||
else:
|
||||
raise ValueError(
|
||||
"Prompt must be string, list of strings, or list of dicts."
|
||||
f"Got {prompt}"
|
||||
)
|
||||
new_request = {
|
||||
"instances": [{"messages": messages}],
|
||||
"parameters": request,
|
||||
}
|
||||
return super(GoogleClient, self).preprocess_request_params(new_request)
|
||||
|
||||
def postprocess_response(self, response: Dict, request: Dict) -> Dict[str, Any]:
|
||||
"""
|
||||
Validate response as dict.
|
||||
|
||||
Assumes response is dict
|
||||
{
|
||||
"candidates": [
|
||||
{
|
||||
"safetyAttributes": {
|
||||
"categories": ["Violent", "Sexual"],
|
||||
"blocked": false,
|
||||
"scores": [0.1, 0.1]
|
||||
},
|
||||
"author": "1",
|
||||
"content": "SELECT * FROM "WWW";"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
Args:
|
||||
response: response
|
||||
request: request
|
||||
|
||||
Return:
|
||||
response as dict
|
||||
"""
|
||||
google_predictions = response.pop("predictions")
|
||||
new_response = {
|
||||
"choices": [
|
||||
{
|
||||
"text": prediction["candidates"][0]["content"],
|
||||
}
|
||||
for prediction in google_predictions
|
||||
]
|
||||
}
|
||||
return super(GoogleClient, self).postprocess_response(new_response, request)
|
||||
Loading…
Reference in New Issue