#!/usr/bin/env python3
# Copyright (c) OpenMMLab. All rights reserved.
"""HuixiangDou binary."""
import argparse
import os
import time
import json
import random
from multiprocessing import Process, Value
import pytoml
import requests
from aiohttp import web
from loguru import logger
from .service import ErrorCode, Worker, llm_serve
def parse_args():
"""Parse args."""
parser = argparse.ArgumentParser(description='Worker.')
parser.add_argument('--work_dir',
type=str,
default='workdir',
help='Working directory.')
parser.add_argument(
'--config_path',
default='config.ini',
type=str,
help='Worker configuration path. Default value is config.ini')
parser.add_argument('--standalone',
action='store_true',
default=False,
help='Auto deploy required Hybrid LLM Service.')
parser.add_argument('--step_by_step',
default='annotate', # 'annotate' or 'sparkle' or 'writting'
help='step by step mode')
args = parser.parse_args()
return args
def check_env(args):
"""Check or create config.ini and logs dir."""
if not os.path.exists('logs'):
os.makedirs('logs')
CONFIG_NAME = 'config.ini'
CONFIG_URL = 'https://raw.githubusercontent.com/InternLM/HuixiangDou/main/config.ini' # noqa E501
if not os.path.exists(CONFIG_NAME):
logger.warning(
f'{CONFIG_NAME} not found, download a template from {CONFIG_URL}.')
try:
response = requests.get(CONFIG_URL, timeout=60)
response.raise_for_status()
with open(CONFIG_NAME, 'wb') as f:
f.write(response.content)
except Exception as e:
logger.error(f'Failed to download file due to {e}')
raise e
if not os.path.exists(args.work_dir):
logger.warning(
f'args.work_dir dir not exist, auto create {args.work_dir}.')
os.makedirs(args.work_dir)
def build_reply_text(reply: str, references: list):
if len(references) < 1:
return reply
ret = reply
for ref in references:
ret += '\n'
ret += ref
return ret
def annotation(assistant, config: dict,k,n=50):
query = 'annotation'
feature_dir = os.path.join(config['feature_store']['work_dir'], 'cluster_features')
samples_json = os.path.join(feature_dir, f'cluster_features_{k}','samples.json')
with open(samples_json, 'r') as f:
samples = json.load(f)
f.close()
new_obj_list = []
for cluster_no in random.sample(samples.keys(), n):
chunk = '\n'.join(samples[cluster_no]['samples'][:10])
code, reply, cluster_no = assistant.annotate_cluster(
cluster_no=cluster_no,
chunk=chunk,
history=[],
groupname='')
references = f"cluster_no: {cluster_no}"
new_obj = {
'cluster_no': cluster_no,
'chunk': chunk,
'annotation': reply
}
new_obj_list.append(new_obj)
logger.info(f'{code}, {query}, {reply}, {references}')
with open(os.path.join(feature_dir, f'cluster_features_{k}', 'annotation.json'), 'w') as f:
json.dump(new_obj_list, f, indent=4, ensure_ascii=False)
f.close()
# def lark_send_only(assistant, fe_config: dict):
# queries = ['what is skin-gut axis?',"什么是肠皮轴？","肠道和皮肤的免疫细胞如何相互影响"]
# for query in queries:
# code, reply, references = assistant.generate(query=query,
# history=[],
# groupname='')
# logger.info(f'{code}, {query}, {reply}, {references}')
# reply_text = build_reply_text(reply=reply, references=references)
# if fe_config['type'] == 'lark' and code == ErrorCode.SUCCESS:
# # send message to lark group
# from .frontend import Lark
# lark = Lark(webhook=fe_config['webhook_url'])
# logger.info(f'send {reply} and {references} to lark group.')
# lark.send_text(msg=reply_text)
# def lark_group_recv_and_send(assistant, fe_config: dict):
# from .frontend import (is_revert_command, revert_from_lark_group,
# send_to_lark_group)
# msg_url = fe_config['webhook_url']
# lark_group_config = fe_config['lark_group']
# sent_msg_ids = []
# while True:
# # fetch a user message
# resp = requests.post(msg_url, timeout=10)
# resp.raise_for_status()
# json_obj = resp.json()
# if len(json_obj) < 1:
# # no user input, sleep
# time.sleep(2)
# continue
# logger.debug(json_obj)
# query = json_obj['content']
# if is_revert_command(query):
# for msg_id in sent_msg_ids:
# error = revert_from_lark_group(msg_id,
# lark_group_config['app_id'],
# lark_group_config['app_secret'])
# if error is not None:
# logger.error(
# f'revert msg_id {msg_id} fail, reason {error}')
# else:
# logger.debug(f'revert msg_id {msg_id}')
# time.sleep(0.5)
# sent_msg_ids = []
# continue
# code, reply, references = assistant.generate(query=query,
# history=[],
# groupname='')
# if code == ErrorCode.SUCCESS:
# json_obj['reply'] = build_reply_text(reply=reply,
# references=references)
# error, msg_id = send_to_lark_group(
# json_obj=json_obj,
# app_id=lark_group_config['app_id'],
# app_secret=lark_group_config['app_secret'])
# if error is not None:
# raise error
# sent_msg_ids.append(msg_id)
# else:
# logger.debug(f'{code} for the query {query}')
# def wechat_personal_run(assistant, fe_config: dict):
# """Call assistant inference."""
# async def api(request):
# input_json = await request.json()
# logger.debug(input_json)
# query = input_json['query']
# if type(query) is dict:
# query = query['content']
# code, reply, references = assistant.generate(query=query,
# history=[],
# groupname='')
# reply_text = build_reply_text(reply=reply, references=references)
# return web.json_response({'code': int(code), 'reply': reply_text})
# bind_port = fe_config['wechat_personal']['bind_port']
# app = web.Application()
# app.add_routes([web.post('/api', api)])
# web.run_app(app, host='0.0.0.0', port=bind_port)
def run():
"""Automatically download config, start llm server and run examples."""
args = parse_args()
check_env(args)
if args.standalone is True:
# hybrid llm serve
server_ready = Value('i', 0)
server_process = Process(target=llm_serve,
args=(args.config_path, server_ready))
server_process.daemon = True
server_process.start()
while True:
if server_ready.value == 0:
logger.info('waiting for server to be ready..')
time.sleep(3)
elif server_ready.value == 1:
break
else:
logger.error('start local LLM server failed, quit.')
raise Exception('local LLM path')
logger.info('Hybrid LLM Server start.')
# query by worker
with open(args.config_path, encoding='utf8') as f:
config = pytoml.load(f)
fe_config = config['frontend']
logger.info('Config loaded.')
assistant = Worker(work_dir=args.work_dir, config_path=args.config_path,language='en')
step = args.step_by_step
if step == 'annotate':
annotation(assistant, config, 500)
annotation(assistant, config, 200)
annotation(assistant, config, 100)
annotation(assistant, config, 50)
annotation(assistant, config, 20,n=20)
annotation(assistant, config, 10,n=10)
elif step == 'sparkle':
pass # TODO
elif step == 'writting':
pass # TODO
else:
logger.info(f'unsupported step_by_step mode {step}, please read `config.ini` description.')
# fe_type = fe_config['type']
# if fe_type == 'lark' or fe_type == 'none':
# lark_send_only(assistant, fe_config)
# elif fe_type == 'lark_group':
# lark_group_recv_and_send(assistant, fe_config)
# elif fe_type == 'wechat_personal':
# wechat_personal_run(assistant, fe_config)
# else:
# logger.info(
# f'unsupported fe_config.type {fe_type}, please read `config.ini` description.' # noqa E501
# )
# server_process.join()
if __name__ == '__main__':
run()