huixiangdou/main.py

#!/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()