大活
This commit is contained in:
538
art-agent/backend/venv/Lib/site-packages/replicate/model.py
Normal file
538
art-agent/backend/venv/Lib/site-packages/replicate/model.py
Normal file
@@ -0,0 +1,538 @@
|
||||
from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, Union, overload
|
||||
|
||||
from typing_extensions import NotRequired, TypedDict, Unpack, deprecated
|
||||
|
||||
from replicate.exceptions import ReplicateException
|
||||
from replicate.helpers import async_encode_json, encode_json
|
||||
from replicate.identifier import ModelVersionIdentifier
|
||||
from replicate.pagination import Page
|
||||
from replicate.prediction import (
|
||||
Prediction,
|
||||
_create_prediction_body,
|
||||
_create_prediction_request_params,
|
||||
_json_to_prediction,
|
||||
)
|
||||
from replicate.resource import Namespace, Resource
|
||||
from replicate.version import Version, Versions
|
||||
|
||||
try:
|
||||
from pydantic import v1 as pydantic # type: ignore
|
||||
except ImportError:
|
||||
import pydantic # type: ignore
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from replicate.client import Client
|
||||
from replicate.prediction import Predictions
|
||||
|
||||
|
||||
class Model(Resource):
|
||||
"""
|
||||
A machine learning model hosted on Replicate.
|
||||
"""
|
||||
|
||||
_client: "Client" = pydantic.PrivateAttr()
|
||||
|
||||
url: str
|
||||
"""
|
||||
The URL of the model.
|
||||
"""
|
||||
|
||||
owner: str
|
||||
"""
|
||||
The owner of the model.
|
||||
"""
|
||||
|
||||
name: str
|
||||
"""
|
||||
The name of the model.
|
||||
"""
|
||||
|
||||
description: Optional[str]
|
||||
"""
|
||||
The description of the model.
|
||||
"""
|
||||
|
||||
visibility: Literal["public", "private"]
|
||||
"""
|
||||
The visibility of the model. Can be 'public' or 'private'.
|
||||
"""
|
||||
|
||||
github_url: Optional[str]
|
||||
"""
|
||||
The GitHub URL of the model.
|
||||
"""
|
||||
|
||||
paper_url: Optional[str]
|
||||
"""
|
||||
The URL of the paper related to the model.
|
||||
"""
|
||||
|
||||
license_url: Optional[str]
|
||||
"""
|
||||
The URL of the license for the model.
|
||||
"""
|
||||
|
||||
run_count: int
|
||||
"""
|
||||
The number of runs of the model.
|
||||
"""
|
||||
|
||||
cover_image_url: Optional[str]
|
||||
"""
|
||||
The URL of the cover image for the model.
|
||||
"""
|
||||
|
||||
default_example: Optional[Prediction]
|
||||
"""
|
||||
The default example of the model.
|
||||
"""
|
||||
|
||||
latest_version: Optional[Version]
|
||||
"""
|
||||
The latest version of the model.
|
||||
"""
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
"""
|
||||
Return the qualified model name, in the format `owner/name`.
|
||||
"""
|
||||
return f"{self.owner}/{self.name}"
|
||||
|
||||
@property
|
||||
@deprecated("Use `model.owner` instead.")
|
||||
def username(self) -> str:
|
||||
"""
|
||||
The name of the user or organization that owns the model.
|
||||
This attribute is deprecated and will be removed in future versions.
|
||||
"""
|
||||
return self.owner
|
||||
|
||||
@username.setter
|
||||
@deprecated("Use `model.owner` instead.")
|
||||
def username(self, value: str) -> None:
|
||||
self.owner = value
|
||||
|
||||
def predict(self, *args, **kwargs) -> None:
|
||||
"""
|
||||
DEPRECATED: Use `replicate.run()` instead.
|
||||
"""
|
||||
|
||||
raise ReplicateException(
|
||||
"The `model.predict()` method has been removed, because it's unstable: if a new version of the model you're using is pushed and its API has changed, your code may break. Use `replicate.run()` instead. See https://github.com/replicate/replicate-python#readme"
|
||||
)
|
||||
|
||||
@property
|
||||
def versions(self) -> Versions:
|
||||
"""
|
||||
Get the versions of this model.
|
||||
"""
|
||||
|
||||
return Versions(client=self._client, model=self)
|
||||
|
||||
def reload(self) -> None:
|
||||
"""
|
||||
Load this object from the server.
|
||||
"""
|
||||
|
||||
obj = self._client.models.get(f"{self.owner}/{self.name}")
|
||||
for name, value in obj.dict().items():
|
||||
setattr(self, name, value)
|
||||
|
||||
|
||||
class Models(Namespace):
|
||||
"""
|
||||
Namespace for operations related to models.
|
||||
"""
|
||||
|
||||
model = Model
|
||||
|
||||
@property
|
||||
def predictions(self) -> "ModelsPredictions":
|
||||
"""
|
||||
Get a namespace for operations related to predictions on a model.
|
||||
"""
|
||||
|
||||
return ModelsPredictions(client=self._client)
|
||||
|
||||
def list(self, cursor: Union[str, "ellipsis", None] = ...) -> Page[Model]: # noqa: F821
|
||||
"""
|
||||
List all public models.
|
||||
|
||||
Parameters:
|
||||
cursor: The cursor to use for pagination. Use the value of `Page.next` or `Page.previous`.
|
||||
Returns:
|
||||
Page[Model]: A page of of models.
|
||||
Raises:
|
||||
ValueError: If `cursor` is `None`.
|
||||
"""
|
||||
|
||||
if cursor is None:
|
||||
raise ValueError("cursor cannot be None")
|
||||
|
||||
resp = self._client._request("GET", "/v1/models" if cursor is ... else cursor)
|
||||
|
||||
obj = resp.json()
|
||||
obj["results"] = [
|
||||
_json_to_model(self._client, result) for result in obj["results"]
|
||||
]
|
||||
|
||||
return Page[Model](**obj)
|
||||
|
||||
async def async_list(
|
||||
self,
|
||||
cursor: Union[str, "ellipsis", None] = ..., # noqa: F821
|
||||
) -> Page[Model]:
|
||||
"""
|
||||
List all public models.
|
||||
|
||||
Parameters:
|
||||
cursor: The cursor to use for pagination. Use the value of `Page.next` or `Page.previous`.
|
||||
Returns:
|
||||
Page[Model]: A page of of models.
|
||||
Raises:
|
||||
ValueError: If `cursor` is `None`.
|
||||
"""
|
||||
|
||||
if cursor is None:
|
||||
raise ValueError("cursor cannot be None")
|
||||
|
||||
resp = await self._client._async_request(
|
||||
"GET", "/v1/models" if cursor is ... else cursor
|
||||
)
|
||||
|
||||
obj = resp.json()
|
||||
obj["results"] = [
|
||||
_json_to_model(self._client, result) for result in obj["results"]
|
||||
]
|
||||
|
||||
return Page[Model](**obj)
|
||||
|
||||
def search(self, query: str) -> Page[Model]:
|
||||
"""
|
||||
Search for public models.
|
||||
|
||||
Parameters:
|
||||
query: The search query.
|
||||
Returns:
|
||||
Page[Model]: A page of models matching the search query.
|
||||
"""
|
||||
resp = self._client._request(
|
||||
"QUERY", "/v1/models", content=query, headers={"Content-Type": "text/plain"}
|
||||
)
|
||||
|
||||
obj = resp.json()
|
||||
obj["results"] = [
|
||||
_json_to_model(self._client, result) for result in obj["results"]
|
||||
]
|
||||
|
||||
return Page[Model](**obj)
|
||||
|
||||
async def async_search(self, query: str) -> Page[Model]:
|
||||
"""
|
||||
Asynchronously search for public models.
|
||||
|
||||
Parameters:
|
||||
query: The search query.
|
||||
Returns:
|
||||
Page[Model]: A page of models matching the search query.
|
||||
"""
|
||||
resp = await self._client._async_request(
|
||||
"QUERY", "/v1/models", content=query, headers={"Content-Type": "text/plain"}
|
||||
)
|
||||
|
||||
obj = resp.json()
|
||||
obj["results"] = [
|
||||
_json_to_model(self._client, result) for result in obj["results"]
|
||||
]
|
||||
|
||||
return Page[Model](**obj)
|
||||
|
||||
@overload
|
||||
def get(self, key: str) -> Model: ...
|
||||
|
||||
@overload
|
||||
def get(self, owner: str, name: str) -> Model: ...
|
||||
|
||||
def get(self, *args, **kwargs) -> Model:
|
||||
"""
|
||||
Get a model by name.
|
||||
"""
|
||||
|
||||
url = _get_model_url(*args, **kwargs)
|
||||
resp = self._client._request("GET", url)
|
||||
|
||||
return _json_to_model(self._client, resp.json())
|
||||
|
||||
@overload
|
||||
async def async_get(self, key: str) -> Model: ...
|
||||
|
||||
@overload
|
||||
async def async_get(self, owner: str, name: str) -> Model: ...
|
||||
|
||||
async def async_get(self, *args, **kwargs) -> Model:
|
||||
"""
|
||||
Get a model by name.
|
||||
|
||||
Args:
|
||||
key: The qualified name of the model, in the format `owner/name`.
|
||||
Returns:
|
||||
The model.
|
||||
"""
|
||||
|
||||
url = _get_model_url(*args, **kwargs)
|
||||
resp = await self._client._async_request("GET", url)
|
||||
|
||||
return _json_to_model(self._client, resp.json())
|
||||
|
||||
@overload
|
||||
def delete(self, key: str) -> bool: ...
|
||||
|
||||
@overload
|
||||
def delete(self, owner: str, name: str) -> bool: ...
|
||||
|
||||
def delete(self, *args, **kwargs) -> bool:
|
||||
"""
|
||||
Delete a model by name.
|
||||
|
||||
Returns:
|
||||
`True` if deletion was successful, otherwise `False`.
|
||||
"""
|
||||
url = _delete_model_url(*args, **kwargs)
|
||||
resp = self._client._request("DELETE", url)
|
||||
return resp.status_code == 204
|
||||
|
||||
@overload
|
||||
async def async_delete(self, key: str) -> bool: ...
|
||||
|
||||
@overload
|
||||
async def async_delete(self, owner: str, name: str) -> bool: ...
|
||||
|
||||
async def async_delete(self, *args, **kwargs) -> bool:
|
||||
"""
|
||||
Asynchronously delete a model by name.
|
||||
|
||||
Returns:
|
||||
`True` if deletion was successful, otherwise `False`.
|
||||
"""
|
||||
url = _delete_model_url(*args, **kwargs)
|
||||
resp = await self._client._async_request("DELETE", url)
|
||||
return resp.status_code == 204
|
||||
|
||||
class CreateModelParams(TypedDict):
|
||||
"""Parameters for creating a model."""
|
||||
|
||||
hardware: str
|
||||
"""The SKU for the hardware used to run the model.
|
||||
|
||||
Possible values can be found by calling `replicate.hardware.list()`."""
|
||||
|
||||
visibility: Literal["public", "private"]
|
||||
"""Whether the model should be public or private."""
|
||||
|
||||
description: NotRequired[str]
|
||||
"""The description of the model."""
|
||||
|
||||
github_url: NotRequired[str]
|
||||
"""A URL for the model's source code on GitHub."""
|
||||
|
||||
paper_url: NotRequired[str]
|
||||
"""A URL for the model's paper."""
|
||||
|
||||
license_url: NotRequired[str]
|
||||
"""A URL for the model's license."""
|
||||
|
||||
cover_image_url: NotRequired[str]
|
||||
"""A URL for the model's cover image."""
|
||||
|
||||
def create(
|
||||
self,
|
||||
owner: str,
|
||||
name: str,
|
||||
**params: Unpack["Models.CreateModelParams"],
|
||||
) -> Model:
|
||||
"""
|
||||
Create a model.
|
||||
"""
|
||||
|
||||
body = _create_model_body(owner, name, **params)
|
||||
resp = self._client._request("POST", "/v1/models", json=body)
|
||||
|
||||
return _json_to_model(self._client, resp.json())
|
||||
|
||||
async def async_create(
|
||||
self, owner: str, name: str, **params: Unpack["Models.CreateModelParams"]
|
||||
) -> Model:
|
||||
"""
|
||||
Create a model.
|
||||
"""
|
||||
|
||||
body = body = _create_model_body(owner, name, **params)
|
||||
resp = await self._client._async_request("POST", "/v1/models", json=body)
|
||||
|
||||
return _json_to_model(self._client, resp.json())
|
||||
|
||||
|
||||
class ModelsPredictions(Namespace):
|
||||
"""
|
||||
Namespace for operations related to predictions in a deployment.
|
||||
"""
|
||||
|
||||
def create(
|
||||
self,
|
||||
model: Union[str, Tuple[str, str], "Model"],
|
||||
input: Dict[str, Any],
|
||||
**params: Unpack["Predictions.CreatePredictionParams"],
|
||||
) -> Prediction:
|
||||
"""
|
||||
Create a new prediction with the deployment.
|
||||
"""
|
||||
|
||||
wait = params.pop("wait", None)
|
||||
file_encoding_strategy = params.pop("file_encoding_strategy", None)
|
||||
|
||||
path = _create_prediction_path_from_model(model)
|
||||
if input is not None:
|
||||
input = encode_json(
|
||||
input,
|
||||
client=self._client,
|
||||
file_encoding_strategy=file_encoding_strategy,
|
||||
)
|
||||
|
||||
body = _create_prediction_body(version=None, input=input, **params)
|
||||
extras = _create_prediction_request_params(wait=wait)
|
||||
resp = self._client._request("POST", path, json=body, **extras)
|
||||
|
||||
return _json_to_prediction(self._client, resp.json())
|
||||
|
||||
async def async_create(
|
||||
self,
|
||||
model: Union[str, Tuple[str, str], "Model"],
|
||||
input: Dict[str, Any],
|
||||
**params: Unpack["Predictions.CreatePredictionParams"],
|
||||
) -> Prediction:
|
||||
"""
|
||||
Create a new prediction with the deployment.
|
||||
"""
|
||||
|
||||
wait = params.pop("wait", None)
|
||||
file_encoding_strategy = params.pop("file_encoding_strategy", None)
|
||||
|
||||
path = _create_prediction_path_from_model(model)
|
||||
|
||||
if input is not None:
|
||||
input = await async_encode_json(
|
||||
input,
|
||||
client=self._client,
|
||||
file_encoding_strategy=file_encoding_strategy,
|
||||
)
|
||||
|
||||
body = _create_prediction_body(version=None, input=input, **params)
|
||||
extras = _create_prediction_request_params(wait=wait)
|
||||
resp = await self._client._async_request("POST", path, json=body, **extras)
|
||||
|
||||
return _json_to_prediction(self._client, resp.json())
|
||||
|
||||
|
||||
def _create_model_body( # pylint: disable=too-many-arguments
|
||||
owner: str,
|
||||
name: str,
|
||||
*,
|
||||
visibility: str,
|
||||
hardware: str,
|
||||
description: Optional[str] = None,
|
||||
github_url: Optional[str] = None,
|
||||
paper_url: Optional[str] = None,
|
||||
license_url: Optional[str] = None,
|
||||
cover_image_url: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
body = {
|
||||
"owner": owner,
|
||||
"name": name,
|
||||
"visibility": visibility,
|
||||
"hardware": hardware,
|
||||
}
|
||||
|
||||
if description is not None:
|
||||
body["description"] = description
|
||||
|
||||
if github_url is not None:
|
||||
body["github_url"] = github_url
|
||||
|
||||
if paper_url is not None:
|
||||
body["paper_url"] = paper_url
|
||||
|
||||
if license_url is not None:
|
||||
body["license_url"] = license_url
|
||||
|
||||
if cover_image_url is not None:
|
||||
body["cover_image_url"] = cover_image_url
|
||||
|
||||
return body
|
||||
|
||||
|
||||
def _get_model_url(*args, **kwargs) -> str:
|
||||
if len(args) > 0 and len(kwargs) > 0:
|
||||
raise ValueError("Cannot mix positional and keyword arguments")
|
||||
|
||||
owner = kwargs.get("owner", None)
|
||||
name = kwargs.get("name", None)
|
||||
key = kwargs.get("key", None)
|
||||
|
||||
if key and (owner or name):
|
||||
raise ValueError(
|
||||
"Must specify exactly one of 'owner' and 'name' or single 'key' in the format 'owner/name'"
|
||||
)
|
||||
|
||||
if args:
|
||||
if len(args) == 1:
|
||||
key = args[0]
|
||||
elif len(args) == 2:
|
||||
owner, name = args
|
||||
else:
|
||||
raise ValueError("Invalid number of arguments")
|
||||
|
||||
if not key:
|
||||
if not (owner and name):
|
||||
raise ValueError(
|
||||
"Both 'owner' and 'name' must be provided if 'key' is not specified."
|
||||
)
|
||||
key = f"{owner}/{name}"
|
||||
|
||||
return f"/v1/models/{key}"
|
||||
|
||||
|
||||
def _delete_model_url(*args, **kwargs) -> str:
|
||||
return _get_model_url(*args, **kwargs)
|
||||
|
||||
|
||||
def _json_to_model(client: "Client", json: Dict[str, Any]) -> Model:
|
||||
model = Model(**json)
|
||||
model._client = client
|
||||
if model.default_example is not None:
|
||||
model.default_example._client = client
|
||||
return model
|
||||
|
||||
|
||||
def _create_prediction_path_from_model(
|
||||
model: Union[str, Tuple[str, str], "Model"],
|
||||
) -> str:
|
||||
owner, name = None, None
|
||||
if isinstance(model, Model):
|
||||
owner, name = model.owner, model.name
|
||||
elif isinstance(model, tuple):
|
||||
owner, name = model[0], model[1]
|
||||
elif isinstance(model, str):
|
||||
owner, name, version_id = ModelVersionIdentifier.parse(model)
|
||||
if version_id is not None:
|
||||
raise ValueError(
|
||||
f"Invalid reference to model version: {model}. Expected model or reference in the format owner/name"
|
||||
)
|
||||
|
||||
if owner is None or name is None:
|
||||
raise ValueError(
|
||||
"model must be a Model, a tuple of (owner, name), or a string in the format 'owner/name'"
|
||||
)
|
||||
|
||||
return f"/v1/models/{owner}/{name}/predictions"
|
||||
Reference in New Issue
Block a user