FilesExpand file tree

 main

/

training.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

90 lines (79 loc) · 2.76 KB

 main

/

training.py

Copy path

Top

File metadata and controls

• Code
• Blame

90 lines (79 loc) · 2.76 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

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

import re

from typing import Any, Dict, List, Optional

from replicate.base_model import BaseModel

from replicate.collection import Collection

from replicate.exceptions import ReplicateException

from replicate.files import upload_file

from replicate.json import encode_json

from replicate.version import Version

class Training(BaseModel):

completed_at: Optional[str]

created_at: Optional[str]

destination: Optional[str]

error: Optional[str]

id: str

input: Optional[Dict[str, Any]]

logs: Optional[str]

output: Optional[Any]

started_at: Optional[str]

status: str

version: Optional[Version]

def cancel(self) -> None:

"""Cancel a running training"""

self._client._request("POST", f"/v1/trainings/{self.id}/cancel")

class TrainingCollection(Collection):

model = Training

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

resp = self._client._request("GET", "/v1/trainings")

# TODO: paginate

trainings = resp.json()["results"]

for training in trainings:

# HACK: resolve this? make it lazy somehow?

del training["version"]

return [self.prepare_model(obj) for obj in trainings]

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

resp = self._client._request(

"GET",

f"/v1/trainings/{id}",

)

obj = resp.json()

# HACK: resolve this? make it lazy somehow?

del obj["version"]

return self.prepare_model(obj)

def create( # type: ignore

self,

version: str,

input: Dict[str, Any],

destination: str,

webhook: Optional[str] = None,

webhook_events_filter: Optional[List[str]] = None,

**kwargs,

) -> Training:

input = encode_json(input, upload_file=upload_file)

body = {

"input": input,

"destination": destination,

}

if webhook is not None:

body["webhook"] = webhook

if webhook_events_filter is not None:

body["webhook_events_filter"] = webhook_events_filter

# Split version in format "username/model_name:version_id"

match = re.match(

r"^(?P<username>[^/]+)/(?P<model_name>[^:]+):(?P<version_id>.+)$", version

)

if not match:

raise ReplicateException(

"version must be in format username/model_name:version_id"

)

username = match.group("username")

model_name = match.group("model_name")

version_id = match.group("version_id")

resp = self._client._request(

"POST",

f"/v1/models/{username}/{model_name}/versions/{version_id}/trainings",

json=body,

)

obj = resp.json()

del obj["version"]

return self.prepare_model(obj)