fix bugs
This commit is contained in:
@@ -6,7 +6,7 @@ from typing import Optional, List
|
|||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from llama_cpp import Llama
|
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
|
from .queue_compute_node import Queue_ComputeNode
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -22,19 +22,26 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
|
|||||||
self.model_name = model_name
|
self.model_name = model_name
|
||||||
self.llm = Llama(model_path=model_path)
|
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:
|
match task.task_type:
|
||||||
case ComputeTaskType.TEXT_EMBEDDING:
|
case ComputeTaskType.TEXT_EMBEDDING:
|
||||||
model_name = task.params["model_name"]
|
model_name = task.params["model_name"]
|
||||||
input = task.params["input"]
|
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}")
|
logger.info(f"local-llama({self.model_path}) response: {embedding}")
|
||||||
|
task.state = ComputeTaskState.DONE
|
||||||
result = ComputeTaskResult()
|
result.result_code = ComputeTaskResultCode.OK
|
||||||
result.set_from_task(task)
|
|
||||||
result.result = embedding
|
result.result = embedding
|
||||||
|
|
||||||
return result
|
return result
|
||||||
@@ -48,19 +55,24 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
|
|||||||
|
|
||||||
logger.info(f"local-llama({self.model_path}) prompts: {prompts}")
|
logger.info(f"local-llama({self.model_path}) prompts: {prompts}")
|
||||||
|
|
||||||
resp = self.llm.create_chat_completion(model=mode_name,
|
try:
|
||||||
messages=prompts,
|
resp = self.llm.create_chat_completion(model=mode_name,
|
||||||
functions=llm_inner_functions, # function has not support?
|
messages=prompts,
|
||||||
max_tokens=max_token_size,
|
functions=llm_inner_functions, # function has not support?
|
||||||
temperature=0.7) # TODO: add temperature to task params?
|
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)}")
|
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"]
|
status_code = resp["choices"][0]["finish_reason"]
|
||||||
|
token_usage = resp["usage"]
|
||||||
match status_code:
|
match status_code:
|
||||||
case "function_call":
|
case "function_call":
|
||||||
task.state = ComputeTaskState.DONE
|
task.state = ComputeTaskState.DONE
|
||||||
@@ -69,11 +81,25 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
|
|||||||
case _:
|
case _:
|
||||||
task.state = ComputeTaskState.ERROR
|
task.state = ComputeTaskState.ERROR
|
||||||
task.error_str = f"The status code was {status_code}."
|
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
|
return None
|
||||||
|
|
||||||
|
result.result_code = ComputeTaskResultCode.OK
|
||||||
result.result_str = resp["choices"][0]["message"]["content"]
|
result.result_str = resp["choices"][0]["message"]["content"]
|
||||||
result.result_message = resp["choices"][0]["message"]
|
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
|
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:
|
async def initial(self) -> bool:
|
||||||
return True
|
return True
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ class Queue_ComputeNode(ComputeNode):
|
|||||||
self.is_start = False
|
self.is_start = False
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
async def execute_task(self, task: ComputeTask) -> ComputeTaskResult:
|
async def execute_task(self, task: ComputeTask, result: ComputeTaskResult):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
async def push_task(self, task: ComputeTask, proiority: int = 0):
|
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):
|
async def _run_task(self, task: ComputeTask):
|
||||||
task.state = ComputeTaskState.RUNNING
|
task.state = ComputeTaskState.RUNNING
|
||||||
|
|
||||||
result = await self.execute_task(task)
|
result = ComputeTaskResult()
|
||||||
if result is not None:
|
result.result_code = ComputeTaskResultCode.ERROR
|
||||||
result.set_from_task(task)
|
result.set_from_task(task)
|
||||||
result.worker_id = self.node_id
|
result.worker_id = self.node_id
|
||||||
else:
|
|
||||||
task.state = ComputeTaskState.ERROR
|
await self.execute_task(task, result)
|
||||||
|
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user