Add embeding compute task

This commit is contained in:
tsukasa
2023-08-31 15:45:02 +08:00
parent 39eb96f9e3
commit febe28eecc
3 changed files with 40 additions and 24 deletions
+1 -1
View File
@@ -123,7 +123,7 @@
</mxGraphModel> </mxGraphModel>
</diagram> </diagram>
<diagram id="kWxfmPxtNxOAf0a73TCG" name="Page-2"> <diagram id="kWxfmPxtNxOAf0a73TCG" name="Page-2">
<mxGraphModel dx="500" dy="864" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="850" pageHeight="1100" math="0" shadow="0"> <mxGraphModel dx="625" dy="1080" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="850" pageHeight="1100" math="0" shadow="0">
<root> <root>
<mxCell id="0"/> <mxCell id="0"/>
<mxCell id="1" parent="0"/> <mxCell id="1" parent="0"/>
+8
View File
@@ -38,6 +38,14 @@ class ComputeTask:
self.params["model_name"] = "gpt-4-0613" self.params["model_name"] = "gpt-4-0613"
self.params["max_token_size"] = max_token_size self.params["max_token_size"] = max_token_size
def set_embeding_params(self, model_name, input, callchain_id = None):
self.task_type = "embeding"
self.create_time = time.time()
self.task_id = uuid.uuid4().hex
self.callchain_id = callchain_id
self.params["model_name"] = model_name
self.params["input"] = input
def display(self) -> str: def display(self) -> str:
return f"ComputeTask: {self.task_id} {self.task_type} {self.state}" return f"ComputeTask: {self.task_id} {self.task_type} {self.state}"
+8
View File
@@ -47,6 +47,12 @@ class OpenAI_ComputeNode(ComputeNode):
def _run_task(self,task:ComputeTask): def _run_task(self,task:ComputeTask):
task.state = ComputeTaskState.RUNNING task.state = ComputeTaskState.RUNNING
# switch tsak type
if task.task_type == "llm_completion":
mode_name = task.params["model_name"]
# max_token_size = task.params["max_token_size"]
prompts = task.params["prompts"]
mode_name = task.params["model_name"] mode_name = task.params["model_name"]
# max_token_size = task.params["max_token_size"] # max_token_size = task.params["max_token_size"]
prompts = task.params["prompts"] prompts = task.params["prompts"]
@@ -71,6 +77,8 @@ class OpenAI_ComputeNode(ComputeNode):
result.result = resp["choices"][0]["message"] result.result = resp["choices"][0]["message"]
return result return result
if task.task_type == "embeding":
pass
def start(self): def start(self):
async def _run_task_loop(): async def _run_task_loop():