gyt/opendan
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix bugs

Browse Source
This commit is contained in:
streetycat
2023-09-28 05:03:06 +00:00
committed by zhangzhen
parent 7b5c0103e8
commit 97d69e43f7
2 changed files with 51 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aios_kernel/local_llama_compute_node.py
+42 -16
View File
@@ -6,7 +6,7 @@ from typing import Optional, List
from pydantic import BaseModel
from llama_cpp import Llama
from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType
from .queue_compute_node import Queue_ComputeNode
logger = logging.getLogger(__name__)
@@ -22,19 +22,26 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
self.model_name = model_name
self.llm = Llama(model_path=model_path)
async def execute_task(self, task: ComputeTask) -> ComputeTaskResult:
async def execute_task(self, task: ComputeTask, result: ComputeTaskResult) -> ComputeTaskResult:
match task.task_type:
case ComputeTaskType.TEXT_EMBEDDING:
model_name = task.params["model_name"]
input = task.params["input"]
logger.info(f"call openai {model_name} input: {input}")
logger.info(f"call local-llama {model_name} input: {input}")
embedding = self.llm.embed(input=input)
try:
embedding = self.llm.embed(input=input)
logger.info(f"local-llama({self.model_path}) response: {embedding}")
except Exception as e:
logger.error(f"call local-llama {model_name} run TEXT_EMBEDDING task error: {e}")
task.state = ComputeTaskState.ERROR
task.error_str = str(e)
result.error_str = str(e)
return result
logger.info(f"local-llama({self.model_path}) response: {resp}")
result = ComputeTaskResult()
result.set_from_task(task)
logger.info(f"local-llama({self.model_path}) response: {embedding}")
task.state = ComputeTaskState.DONE
result.result_code = ComputeTaskResultCode.OK
result.result = embedding
return result
@@ -48,19 +55,24 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
logger.info(f"local-llama({self.model_path}) prompts: {prompts}")
resp = self.llm.create_chat_completion(model=mode_name,
messages=prompts,
functions=llm_inner_functions, # function has not support?
max_tokens=max_token_size,
temperature=0.7) # TODO: add temperature to task params?
try:
resp = self.llm.create_chat_completion(model=mode_name,
messages=prompts,
functions=llm_inner_functions, # function has not support?
max_tokens=max_token_size,
temperature=0.7) # TODO: add temperature to task params?
except Exception as e:
logger.error(f"local-llama({self.model_path}) run LLM_COMPLETION task error: {e}")
task.state = ComputeTaskState.ERROR
task.error_str = str(e)
result.error_str = str(e)
return result
logger.info(f"local-llama({self.model_path}) response: {json.dumps(resp, indent=4)}")
result = ComputeTaskResult()
result.set_from_task(task)
status_code = resp["choices"][0]["finish_reason"]
token_usage = resp["usage"]
match status_code:
case "function_call":
task.state = ComputeTaskState.DONE
@@ -69,11 +81,25 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
case _:
task.state = ComputeTaskState.ERROR
task.error_str = f"The status code was {status_code}."
result.error_str = f"The status code was {status_code}."
result.result_code = ComputeTaskResultCode.ERROR
return None
result.result_code = ComputeTaskResultCode.OK
result.result_str = resp["choices"][0]["message"]["content"]
result.result_message = resp["choices"][0]["message"]
if token_usage:
result.result_refers["token_usage"] = token_usage
logger.info(f"local-llama({self.model_path}) success response: {result.result_str}")
return result
case _:
task.state = ComputeTaskState.ERROR
result.result_code = ComputeTaskResultCode.ERROR
task.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
result.error_str = f"ComputeTask's TaskType : {task.task_type} not support!"
return None
async def initial(self) -> bool:
return True
src/aios_kernel/queue_compute_node.py
+7 -7
View File
@@ -16,7 +16,7 @@ class Queue_ComputeNode(ComputeNode):
self.is_start = False
@abstractmethod
async def execute_task(self, task: ComputeTask) -> ComputeTaskResult:
async def execute_task(self, task: ComputeTask, result: ComputeTaskResult):
pass
async def push_task(self, task: ComputeTask, proiority: int = 0):
@@ -29,12 +29,12 @@ class Queue_ComputeNode(ComputeNode):
async def _run_task(self, task: ComputeTask):
task.state = ComputeTaskState.RUNNING
result = await self.execute_task(task)
if result is not None:
result.set_from_task(task)
result.worker_id = self.node_id
else:
task.state = ComputeTaskState.ERROR
result = ComputeTaskResult()
result.result_code = ComputeTaskResultCode.ERROR
result.set_from_task(task)
result.worker_id = self.node_id
await self.execute_task(task, result)
return result