FilesExpand file tree

 main

/

model.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

43 lines (32 loc) · 1.49 KB

 main

/

model.py

Copy path

Top

File metadata and controls

• Code
• Blame

43 lines (32 loc) · 1.49 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

from typing import Dict, List, Union

from replicate.base_model import BaseModel

from replicate.collection import Collection

from replicate.exceptions import ReplicateException

from replicate.version import VersionCollection

class Model(BaseModel):

username: str

name: str

def predict(self, *args, **kwargs) -> None:

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 `version.predict()` instead. See https://github.com/replicate/replicate-python#readme"

)

@property

def versions(self) -> VersionCollection:

return VersionCollection(client=self._client, model=self)

class ModelCollection(Collection):

model = Model

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

raise NotImplementedError()

def get(self, name: str) -> Model:

# TODO: fetch model from server

# TODO: support permanent IDs

username, name = name.split("/")

return self.prepare_model({"username": username, "name": name})

def create(self, **kwargs) -> Model:

raise NotImplementedError()

def prepare_model(self, attrs: Union[Model, Dict]) -> Model:

if isinstance(attrs, BaseModel):

attrs.id = f"{attrs.username}/{attrs.name}"

elif isinstance(attrs, dict):

attrs["id"] = f"{attrs['username']}/{attrs['name']}"

return super().prepare_model(attrs)