streetycat
118 lines
5.1 KiB
Python
118 lines
5.1 KiB
Python
|
|
import json
|
|
import logging
|
|
import requests
|
|
from typing import Optional, List
|
|
from pydantic import BaseModel
|
|
from llama_cpp import Llama
|
|
|
|
from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType
|
|
from .queue_compute_node import Queue_ComputeNode
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
"""
|
|
This is a custom implementation, it should be redesigned.
|
|
"""
|
|
|
|
class LocalLlama_ComputeNode(Queue_ComputeNode):
|
|
def __init__(self, model_path: str, model_name: str):
|
|
super().__init__()
|
|
self.model_path = model_path
|
|
self.model_name = model_name
|
|
self.llm = Llama(model_path=model_path)
|
|
|
|
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 local-llama {model_name} 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: {embedding}")
|
|
task.state = ComputeTaskState.DONE
|
|
result.result_code = ComputeTaskResultCode.OK
|
|
result.result = embedding
|
|
|
|
return result
|
|
case ComputeTaskType.LLM_COMPLETION:
|
|
mode_name = task.params["model_name"]
|
|
prompts = task.params["prompts"]
|
|
max_token_size = task.params.get("max_token_size")
|
|
llm_inner_functions = task.params.get("inner_functions")
|
|
if max_token_size is None:
|
|
max_token_size = 4000
|
|
|
|
logger.info(f"local-llama({self.model_path}) prompts: {prompts}")
|
|
|
|
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)}")
|
|
|
|
status_code = resp["choices"][0]["finish_reason"]
|
|
token_usage = resp["usage"]
|
|
|
|
match status_code:
|
|
case "function_call":
|
|
task.state = ComputeTaskState.DONE
|
|
case "stop":
|
|
task.state = ComputeTaskState.DONE
|
|
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
|
|
|
|
def display(self) -> str:
|
|
return f"LocalLlama_ComputeNode: {self.node_id}"
|
|
|
|
def get_capacity(self):
|
|
pass
|
|
|
|
def is_support(self, task: ComputeTask) -> bool:
|
|
return (task.task_type == ComputeTaskType.TEXT_EMBEDDING or task.task_type == ComputeTaskType.LLM_COMPLETION) and (not task.params["model_name"] or task.params["model_name"] == self.model_name)
|
|
|
|
def is_local(self) -> bool:
|
|
return True
|