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

rebase to main

Browse Source
This commit is contained in:
streetycat
2023-09-28 09:09:30 +00:00
parent 4f6b04fe48
commit 193c627bdb
4 changed files with 255 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aios_kernel/local_llama_compute_node.py
+106 -59
View File
@@ -4,10 +4,10 @@ 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
from .storage import AIStorage,UserConfig
logger = logging.getLogger(__name__)
@@ -16,59 +16,117 @@ This is a custom implementation, it should be redesigned.
"""
class LocalLlama_ComputeNode(Queue_ComputeNode):
def __init__(self, model_path: str, model_name: str):
def __init__(self, url: str, model_name: str):
super().__init__()
self.model_path = model_path
self.url = url
self.model_name = model_name
self.llm = Llama(model_path=model_path)
async def execute_task(self, task: ComputeTask, result: ComputeTaskResult) -> ComputeTaskResult:
async def execute_task(self, task: ComputeTask, result: 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}")
logger.info(f"call local-llama ({self.url}, {self.model_name}) {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
self.embedding(input, result)
logger.info(f"local-llama({self.model_path}) response: {embedding}")
task.state = ComputeTaskState.DONE
result.result_code = ComputeTaskResultCode.OK
result.result = embedding
if result.result_code == ComputeTaskResultCode.OK:
task.state = ComputeTaskState.DONE
else:
task.state = ComputeTaskState.ERROR
task.error_str = result.error_str
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}")
logger.info(f"local-llama({self.url}, {self.model_name}) 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}")
self.completion(task, result)
if result.result_code == ComputeTaskResultCode.OK:
task.state = ComputeTaskState.DONE
else:
task.state = ComputeTaskState.ERROR
task.error_str = str(e)
result.error_str = str(e)
return result
task.error_str = result.error_str
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"local-llama: {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
def embedding(self, input: str, result: ComputeTaskResult):
body = {
"input": input
}
try:
response = requests.post(self.url + "/v1/embeddings", json = body, verify=False, headers={"Content-Type": "application/json"})
response.close()
logger.info(f"local-llama({self.url}, {self.model_name}) task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
if response.status_code == 200:
resp = response.json()
result.result = resp["data"][0]["embedding"]
elif response.status_code == 422:
resp = response.json()
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = "http request failed: " + str(resp["detail"][0]["msg"])
else:
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = "http request failed: " + str(response.status_code)
except Exception as e:
logger.error(f"call local-llama({self.url}, {self.model_name}) run TEXT_EMBEDDING task error: {e}")
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = str(e)
return result
def completion(self, task: ComputeTask, result: ComputeTaskResult):
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 = max_token_size
logger.info(f"local-llama({self.model_path}) response: {json.dumps(resp, indent=4)}")
body = {
"messages": [],
"max_tokens": 4000
}
for prompt in prompts:
body["messages"].append({
"role": prompt["role"],
"content": prompt["content"]
})
try:
response = requests.post(self.url + "/v1/chat/completions", json = body, verify=False, headers={"Content-Type": "application/json"})
response.close()
logger.info(f"local-llama({self.url}, {self.model_name}) task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
if response.status_code == 200:
resp = response.json()
status_code = resp["choices"][0]["finish_reason"]
token_usage = resp["usage"]
@@ -91,27 +149,16 @@ class LocalLlama_ComputeNode(Queue_ComputeNode):
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
logger.info(f"local-llama({self.url}, {self.model_name}) success response: {result.result_str}")
elif response.status_code == 422:
resp = response.json()
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
result.error_str = "http request failed: " + str(resp["detail"][0]["msg"])
else:
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = "http request failed: " + str(response.status_code)
except Exception as e:
logger.error(f"call local-llama({self.url}, {self.model_name}) run LLM_COMPLETION task error: {e}")
result.result_code = ComputeTaskResultCode.ERROR
result.error_str = str(e)
return result