FilesExpand file tree

 main

/

version.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

64 lines (53 loc) · 2.03 KB

 main

/

version.py

Copy path

Top

File metadata and controls

• Code
• Blame

64 lines (53 loc) · 2.03 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

import datetime

from typing import Any, Iterator, List, Union

from replicate.base_model import BaseModel

from replicate.collection import Collection

from replicate.exceptions import ModelError

from replicate.schema import make_schema_backwards_compatible

class Version(BaseModel):

id: str

created_at: datetime.datetime

cog_version: str

openapi_schema: Any

def predict(self, **kwargs) -> Union[Any, Iterator[Any]]:

# TODO: support args

prediction = self._client.predictions.create(version=self, input=kwargs)

# Return an iterator of the output

# FIXME: might just be a list, not an iterator. I wonder if we should differentiate?

schema = self.get_transformed_schema()

output = schema["components"]["schemas"]["Output"]

if (

output.get("type") == "array"

and output.get("x-cog-array-type") == "iterator"

):

return prediction.output_iterator()

prediction.wait()

if prediction.status == "failed":

raise ModelError(prediction.error)

return prediction.output

def get_transformed_schema(self):

schema = self.openapi_schema

schema = make_schema_backwards_compatible(schema, self.cog_version)

return schema

class VersionCollection(Collection):

model = Version

def __init__(self, client, model):

super().__init__(client=client)

self._model = model

# doesn't exist yet

def get(self, id: str) -> Version:

"""

Get a specific version.

"""

resp = self._client._request(

"GET", f"/v1/models/{self._model.username}/{self._model.name}/versions/{id}"

)

return self.prepare_model(resp.json())

def list(self) -> List[Version]:

"""

Return a list of all versions for a model.

"""

resp = self._client._request(

"GET", f"/v1/models/{self._model.username}/{self._model.name}/versions"

)

return [self.prepare_model(obj) for obj in resp.json()["results"]]