-
Notifications
You must be signed in to change notification settings - Fork 131
/
compute_task.py
122 lines (102 loc) · 3.89 KB
/
compute_task.py
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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
from enum import Enum
import uuid
import time
from typing import Union
from knowledge import ObjectID
class ComputeTaskResultCode(Enum):
OK = 0
TIMEOUT = 1
NO_WORKER = 2
ERROR = 3
class ComputeTaskState(Enum):
DONE = 0
INIT = 1
RUNNING = 2
ERROR = 3
PENDING = 4
class ComputeTaskType(Enum):
NONE = "None"
LLM_COMPLETION = "llm_completion"
TEXT_2_IMAGE = "text_2_image"
IMAGE_2_IMAGE = "image_2_image"
VOICE_2_TEXT = "voice_2_text"
TEXT_2_VOICE = "text_2_voice"
TEXT_EMBEDDING ="text_embedding"
IMAGE_EMBEDDING ="image_embedding"
class ComputeTask:
def __init__(self) -> None:
self.task_type = ComputeTaskType.NONE
self.create_time = None
self.task_id: str = None
self.callchain_id: str = None
self.params: dict = {}
self.refers: dict = None
self.pading_data: bytearray = None
self.state = ComputeTaskState.INIT
self.result = None
self.error_str = None
def set_llm_params(self, prompts, model_name, max_token_size, inner_functions = None, callchain_id=None):
self.task_type = ComputeTaskType.LLM_COMPLETION
self.create_time = time.time()
self.task_id = uuid.uuid4().hex
self.callchain_id = callchain_id
self.params["prompts"] = prompts.to_message_list()
if model_name is not None:
self.params["model_name"] = model_name
else:
self.params["model_name"] = "gpt-4-0613"
if max_token_size is None:
self.params["max_token_size"] = 4000
else:
self.params["max_token_size"] = max_token_size
if inner_functions is not None:
self.params["inner_functions"] = inner_functions
def set_text_embedding_params(self, input: str, model_name=None, callchain_id = None):
self.task_type = ComputeTaskType.TEXT_EMBEDDING
self.create_time = time.time()
self.task_id = uuid.uuid4().hex
self.callchain_id = callchain_id
if model_name is not None:
self.params["model_name"] = model_name
else:
self.params["model_name"] = "text-embedding-ada-002"
self.params["input"] = input
def set_image_embedding_params(self, input = Union[ObjectID, bytes], model_name=None, callchain_id = None):
self.task_type = ComputeTaskType.IMAGE_EMBEDDING
self.create_time = time.time()
self.task_id = uuid.uuid4().hex
self.callchain_id = callchain_id
if model_name is not None:
self.params["model_name"] = model_name
else:
self.params["model_name"] = None
self.params["input"] = input
def set_text_2_image_params(self, prompt: str, model_name, negative_prompt="", callchain_id=None):
self.task_type = ComputeTaskType.TEXT_2_IMAGE
self.create_time = time.time()
self.task_id = uuid.uuid4().hex
self.callchain_id = callchain_id
self.params["prompt"] = prompt
self.params["negative_prompt"] = negative_prompt
if model_name is not None:
self.params["model_name"] = model_name
else:
self.params["model_name"] = "v1-5-pruned-emaonly"
def display(self) -> str:
return f"ComputeTask: {self.task_id} {self.task_type} {self.state}"
class ComputeTaskResult:
def __init__(self) -> None:
self.create_time = None
self.task_id: str = None
self.callchain_id: str = None
self.worker_id: str = None
self.error_str : str = None
self.result_code: int = 0
self.result_str: str = None # easy to use,can read from result
self.result : dict = {}
self.result_refers: dict = {}
self.pading_data: bytearray = None
def set_from_task(self, task: ComputeTask):
self.task_id = task.task_id
self.callchain_id = task.callchain_id
task.result = self