Create awesome_chat.py

awesome_chat.py

@@ -0,0 +1,939 @@
+import base64
+import copy
+import datetime
+from io import BytesIO
+import io
+import os
+import random
+import time
+import traceback
+import uuid
+import requests
+import re
+import json
+import logging
+import argparse
+import yaml
+from PIL import Image, ImageDraw
+from diffusers.utils import load_image
+from pydub import AudioSegment
+import threading
+from queue import Queue
+from get_token_ids import get_token_ids_for_task_parsing, get_token_ids_for_choose_model, count_tokens, get_max_context_length
+from huggingface_hub.inference_api import InferenceApi
+from huggingface_hub.inference_api import ALL_TASKS
+from models_server import models, status
+from functools import partial
+from huggingface_hub import Repository
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--config", type=str, default="config.yaml.dev")
+parser.add_argument("--mode", type=str, default="cli")
+args = parser.parse_args()
+
+if __name__ != "__main__":
+    args.config = "config.gradio.yaml"
+
+config = yaml.load(open(args.config, "r"), Loader=yaml.FullLoader)
+
+if not os.path.exists("logs"):
+    os.mkdir("logs")
+
+now = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+
+DATASET_REPO_URL = "https://huggingface.co/datasets/deepcode-ai/HuggingSpace_logs"
+LOG_HF_TOKEN = os.environ.get("LOG_HF_TOKEN")
+if LOG_HF_TOKEN:
+    repo = Repository(
+        local_dir="logs", clone_from=DATASET_REPO_URL, use_auth_token=LOG_HF_TOKEN
+    )
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logger.handlers = []
+logger.propagate = False
+
+handler = logging.StreamHandler()
+formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+handler.setFormatter(formatter)
+if config["debug"]:
+    handler.setLevel(logging.DEBUG)
+logger.addHandler(handler)
+
+log_file = config["log_file"]
+if log_file:
+    log_file = log_file.replace("TIMESTAMP", now)
+    filehandler = logging.FileHandler(log_file)
+    filehandler.setLevel(logging.DEBUG)
+    filehandler.setFormatter(formatter)
+    logger.addHandler(filehandler)
+
+LLM = config["model"]
+use_completion = config["use_completion"]
+
+# consistent: wrong msra model name 
+LLM_encoding = LLM
+if LLM == "gpt-3.5-turbo":
+    LLM_encoding = "text-davinci-003"
+task_parsing_highlight_ids = get_token_ids_for_task_parsing(LLM_encoding)
+choose_model_highlight_ids = get_token_ids_for_choose_model(LLM_encoding)
+
+# ENDPOINT	MODEL NAME	
+# /v1/chat/completions	gpt-4, gpt-4-0314, gpt-4-32k, gpt-4-32k-0314, gpt-3.5-turbo, gpt-3.5-turbo-0301	
+# /v1/completions	text-davinci-003, text-davinci-002, text-curie-001, text-babbage-001, text-ada-001, davinci, curie, babbage, ada
+
+if use_completion:
+    api_name = "completions"
+else:
+    api_name = "chat/completions"
+
+if not config["dev"]:
+    if not config["openai"]["key"].startswith("sk-") and not config["openai"]["key"]=="gradio":
+        raise ValueError("Incrorrect OpenAI key. Please check your config.yaml file.")
+    OPENAI_KEY = config["openai"]["key"]
+    endpoint = f"https://api.openai.com/v1/{api_name}"
+    if OPENAI_KEY.startswith("sk-"):
+        HEADER = {
+            "Authorization": f"Bearer {OPENAI_KEY}"
+        }
+    else:
+        HEADER = None
+else:
+    endpoint = f"{config['local']['endpoint']}/v1/{api_name}"
+    HEADER = None
+
+PROXY = None
+if config["proxy"]:
+    PROXY = {
+        "https": config["proxy"],
+    }
+
+inference_mode = config["inference_mode"]
+
+parse_task_demos_or_presteps = open(config["demos_or_presteps"]["parse_task"], "r").read()
+choose_model_demos_or_presteps = open(config["demos_or_presteps"]["choose_model"], "r").read()
+response_results_demos_or_presteps = open(config["demos_or_presteps"]["response_results"], "r").read()
+
+parse_task_prompt = config["prompt"]["parse_task"]
+choose_model_prompt = config["prompt"]["choose_model"]
+response_results_prompt = config["prompt"]["response_results"]
+
+parse_task_tprompt = config["tprompt"]["parse_task"]
+choose_model_tprompt = config["tprompt"]["choose_model"]
+response_results_tprompt = config["tprompt"]["response_results"]
+
+MODELS = [json.loads(line) for line in open("data/p0_models.jsonl", "r").readlines()]
+MODELS_MAP = {}
+for model in MODELS:
+    tag = model["task"]
+    if tag not in MODELS_MAP:
+        MODELS_MAP[tag] = []
+    MODELS_MAP[tag].append(model)
+METADATAS = {}
+for model in MODELS:
+    METADATAS[model["id"]] = model
+
+def convert_chat_to_completion(data):
+    messages = data.pop('messages', [])
+    tprompt = ""
+    if messages[0]['role'] == "system":
+        tprompt = messages[0]['content']
+        messages = messages[1:]
+    final_prompt = ""
+    for message in messages:
+        if message['role'] == "user":
+            final_prompt += ("<im_start>"+ "user" + "\n" + message['content'] + "<im_end>\n")
+        elif message['role'] == "assistant":
+            final_prompt += ("<im_start>"+ "assistant" + "\n" + message['content'] + "<im_end>\n")
+        else:
+            final_prompt += ("<im_start>"+ "system" + "\n" + message['content'] + "<im_end>\n")
+    final_prompt = tprompt + final_prompt
+    final_prompt = final_prompt + "<im_start>assistant"
+    data["prompt"] = final_prompt
+    data['stop'] = data.get('stop', ["<im_end>"])
+    data['max_tokens'] = data.get('max_tokens', max(get_max_context_length(LLM) - count_tokens(LLM_encoding, final_prompt), 1))
+    return data
+
+def send_request(data):
+    global HEADER
+    openaikey = data.pop("openaikey")
+    if use_completion:
+        data = convert_chat_to_completion(data)
+    if openaikey and openaikey.startswith("sk-"):
+        HEADER = {
+            "Authorization": f"Bearer {openaikey}"
+        }
+    
+    response = requests.post(endpoint, json=data, headers=HEADER, proxies=PROXY)
+    logger.debug(response.text.strip())
+    if "choices" not in response.json():
+        return response.json()
+    if use_completion:
+        return response.json()["choices"][0]["text"].strip()
+    else:
+        return response.json()["choices"][0]["message"]["content"].strip()
+
+def replace_slot(text, entries):
+    for key, value in entries.items():
+        if not isinstance(value, str):
+            value = str(value)
+        text = text.replace("{{" + key +"}}", value.replace('"', "'").replace('\n', "").replace('\\', '\\\\'))
+    return text
+
+def find_json(s):
+    s = s.replace("\'", "\"")
+    start = s.find("{")
+    end = s.rfind("}")
+    res = s[start:end+1]
+    res = res.replace("\n", "")
+    return res
+
+def field_extract(s, field):
+    try:
+        field_rep = re.compile(f'{field}.*?:.*?"(.*?)"', re.IGNORECASE)
+        extracted = field_rep.search(s).group(1).replace("\"", "\'")
+    except:
+        field_rep = re.compile(f'{field}:\ *"(.*?)"', re.IGNORECASE)
+        extracted = field_rep.search(s).group(1).replace("\"", "\'")
+    return extracted
+
+def get_id_reason(choose_str):
+    reason = field_extract(choose_str, "reason")
+    id = field_extract(choose_str, "id")
+    choose = {"id": id, "reason": reason}
+    return id.strip(), reason.strip(), choose
+
+def record_case(success, **args):
+    if not success:
+        return
+    f = open(f"logs/log_success_{now}.jsonl", "a")
+    log = args
+    f.write(json.dumps(log) + "\n")
+    f.close()
+    if LOG_HF_TOKEN:
+        commit_url = repo.push_to_hub(blocking=False)
+
+def image_to_bytes(img_url):
+    img_byte = io.BytesIO()
+    type = img_url.split(".")[-1]
+    load_image(img_url).save(img_byte, format="png")
+    img_data = img_byte.getvalue()
+    return img_data
+
+def resource_has_dep(command):
+    args = command["args"]
+    for _, v in args.items():
+        if "<GENERATED>" in v:
+            return True
+    return False
+
+def fix_dep(tasks):
+    for task in tasks:
+        args = task["args"]
+        task["dep"] = []
+        for k, v in args.items():
+            if "<GENERATED>" in v:
+                dep_task_id = int(v.split("-")[1])
+                if dep_task_id not in task["dep"]:
+                    task["dep"].append(dep_task_id)
+        if len(task["dep"]) == 0:
+            task["dep"] = [-1]
+    return tasks
+
+def unfold(tasks):
+    flag_unfold_task = False
+    try:
+        for task in tasks:
+            for key, value in task["args"].items():
+                if "<GENERATED>" in value:
+                    generated_items = value.split(",")
+                    if len(generated_items) > 1:
+                        flag_unfold_task = True
+                        for item in generated_items:
+                            new_task = copy.deepcopy(task)
+                            dep_task_id = int(item.split("-")[1])
+                            new_task["dep"] = [dep_task_id]
+                            new_task["args"][key] = item
+                            tasks.append(new_task)
+                        tasks.remove(task)
+    except Exception as e:
+        print(e)
+        traceback.print_exc()
+        logger.debug("unfold task failed.")
+
+    if flag_unfold_task:
+        logger.debug(f"unfold tasks: {tasks}")
+        
+    return tasks
+
+def chitchat(messages, openaikey=None):
+    data = {
+        "model": LLM,
+        "messages": messages,
+        "openaikey": openaikey
+    }
+    return send_request(data)
+
+def parse_task(context, input, openaikey=None):
+    demos_or_presteps = parse_task_demos_or_presteps
+    messages = json.loads(demos_or_presteps)
+    for message in messages:
+        if not isinstance(message["content"], str):
+            message["content"] = json.dumps(message["content"], ensure_ascii=False)
+    messages.insert(0, {"role": "system", "content": parse_task_tprompt})
+
+    # cut chat logs
+    start = 0
+    while start <= len(context):
+        history = context[start:]
+        prompt = replace_slot(parse_task_prompt, {
+            "input": input,
+            "context": history 
+        })
+        messages.append({"role": "user", "content": prompt})
+        history_text = "<im_end>\nuser<im_start>".join([m["content"] for m in messages])
+        num = count_tokens(LLM_encoding, history_text)
+        if get_max_context_length(LLM) - num > 800:
+            break
+        messages.pop()
+        start += 2
+    
+    logger.debug(messages)
+    data = {
+        "model": LLM,
+        "messages": messages,
+        "temperature": 0,
+        "logit_bias": {item: config["logit_bias"]["parse_task"] for item in task_parsing_highlight_ids},
+        "openaikey": openaikey
+    }
+    return send_request(data)
+
+def choose_model(input, task, metas, openaikey = None):
+    prompt = replace_slot(choose_model_prompt, {
+        "input": input,
+        "task": task,
+        "metas": metas,
+    })
+    demos_or_presteps = replace_slot(choose_model_demos_or_presteps, {
+        "input": input,
+        "task": task,
+        "metas": metas
+    })
+    messages = json.loads(demos_or_presteps)
+    messages.insert(0, {"role": "system", "content": choose_model_tprompt})
+    messages.append({"role": "user", "content": prompt})
+    logger.debug(messages)
+    data = {
+        "model": LLM,
+        "messages": messages,
+        "temperature": 0,
+        "logit_bias": {item: config["logit_bias"]["choose_model"] for item in choose_model_highlight_ids}, # 5
+        "openaikey": openaikey
+    }
+    return send_request(data)
+
+
+def response_results(input, results, openaikey=None):
+    results = [v for k, v in sorted(results.items(), key=lambda item: item[0])]
+    prompt = replace_slot(response_results_prompt, {
+        "input": input,
+    })
+    demos_or_presteps = replace_slot(response_results_demos_or_presteps, {
+        "input": input,
+        "processes": results
+    })
+    messages = json.loads(demos_or_presteps, strict=False)
+    messages.insert(0, {"role": "system", "content": response_results_tprompt})
+    messages.append({"role": "user", "content": prompt})
+    logger.debug(messages)
+    data = {
+        "model": LLM,
+        "messages": messages,
+        "temperature": 0,
+        "openaikey": openaikey
+    }
+    return send_request(data)
+
+def huggingface_model_inference(model_id, data, task, huggingfacetoken=None):
+    if huggingfacetoken is None:
+        HUGGINGFACE_HEADERS = {}
+    else:
+        HUGGINGFACE_HEADERS = {
+            "Authorization": f"Bearer {huggingfacetoken}",
+    }
+    task_url = f"https://api-inference.huggingface.co/models/{model_id}" # InferenceApi does not yet support some tasks
+    inference = InferenceApi(repo_id=model_id, token=huggingfacetoken)
+    
+    # NLP tasks
+    if task == "question-answering":
+        inputs = {"question": data["text"], "context": (data["context"] if "context" in data else "" )}
+        result = inference(inputs)
+    if task == "sentence-similarity":
+        inputs = {"source_sentence": data["text1"], "target_sentence": data["text2"]}
+        result = inference(inputs)
+    if task in ["text-classification",  "token-classification", "text2text-generation", "summarization", "translation", "conversational", "text-generation"]:
+        inputs = data["text"]
+        result = inference(inputs)
+    
+    # CV tasks
+    if task == "visual-question-answering" or task == "document-question-answering":
+        img_url = data["image"]
+        text = data["text"]
+        img_data = image_to_bytes(img_url)
+        img_base64 = base64.b64encode(img_data).decode("utf-8")
+        json_data = {}
+        json_data["inputs"] = {}
+        json_data["inputs"]["question"] = text
+        json_data["inputs"]["image"] = img_base64
+        result = requests.post(task_url, headers=HUGGINGFACE_HEADERS, json=json_data).json()
+        # result = inference(inputs) # not support
+
+    if task == "image-to-image":
+        img_url = data["image"]
+        img_data = image_to_bytes(img_url)
+        # result = inference(data=img_data) # not support
+        HUGGINGFACE_HEADERS["Content-Length"] = str(len(img_data))
+        r = requests.post(task_url, headers=HUGGINGFACE_HEADERS, data=img_data)
+        result = r.json()
+        if "path" in result:
+            result["generated image"] = result.pop("path")
+    
+    if task == "text-to-image":
+        inputs = data["text"]
+        img = inference(inputs)
+        name = str(uuid.uuid4())[:4]
+        img.save(f"public/images/{name}.png")
+        result = {}
+        result["generated image"] = f"/images/{name}.png"
+
+    if task == "image-segmentation":
+        img_url = data["image"]
+        img_data = image_to_bytes(img_url)
+        image = Image.open(BytesIO(img_data))
+        predicted = inference(data=img_data)
+        colors = []
+        for i in range(len(predicted)):
+            colors.append((random.randint(100, 255), random.randint(100, 255), random.randint(100, 255), 155))
+        for i, pred in enumerate(predicted):
+            label = pred["label"]
+            mask = pred.pop("mask").encode("utf-8")
+            mask = base64.b64decode(mask)
+            mask = Image.open(BytesIO(mask), mode='r')
+            mask = mask.convert('L')
+
+            layer = Image.new('RGBA', mask.size, colors[i])
+            image.paste(layer, (0, 0), mask)
+        name = str(uuid.uuid4())[:4]
+        image.save(f"public/images/{name}.jpg")
+        result = {}
+        result["generated image with segmentation mask"] = f"/images/{name}.jpg"
+        result["predicted"] = predicted
+
+    if task == "object-detection":
+        img_url = data["image"]
+        img_data = image_to_bytes(img_url)
+        predicted = inference(data=img_data)
+        image = Image.open(BytesIO(img_data))
+        draw = ImageDraw.Draw(image)
+        labels = list(item['label'] for item in predicted)
+        color_map = {}
+        for label in labels:
+            if label not in color_map:
+                color_map[label] = (random.randint(0, 255), random.randint(0, 100), random.randint(0, 255))
+        for label in predicted:
+            box = label["box"]
+            draw.rectangle(((box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])), outline=color_map[label["label"]], width=2)
+            draw.text((box["xmin"]+5, box["ymin"]-15), label["label"], fill=color_map[label["label"]])
+        name = str(uuid.uuid4())[:4]
+        image.save(f"public/images/{name}.jpg")
+        result = {}
+        result["generated image with predicted box"] = f"/images/{name}.jpg"
+        result["predicted"] = predicted
+
+    if task in ["image-classification"]:
+        img_url = data["image"]
+        img_data = image_to_bytes(img_url)
+        result = inference(data=img_data)
+ 
+    if task == "image-to-text":
+        img_url = data["image"]
+        img_data = image_to_bytes(img_url)
+        HUGGINGFACE_HEADERS["Content-Length"] = str(len(img_data))
+        r = requests.post(task_url, headers=HUGGINGFACE_HEADERS, data=img_data)
+        result = {}
+        if "generated_text" in r.json()[0]:
+            result["generated text"] = r.json()[0].pop("generated_text")
+    
+    # AUDIO tasks
+    if task == "text-to-speech":
+        inputs = data["text"]
+        response = inference(inputs, raw_response=True)
+        # response = requests.post(task_url, headers=HUGGINGFACE_HEADERS, json={"inputs": text})
+        name = str(uuid.uuid4())[:4]
+        with open(f"public/audios/{name}.flac", "wb") as f:
+            f.write(response.content)
+        result = {"generated audio": f"/audios/{name}.flac"}
+    if task in ["automatic-speech-recognition", "audio-to-audio", "audio-classification"]:
+        audio_url = data["audio"]
+        audio_data = requests.get(audio_url, timeout=10).content
+        response = inference(data=audio_data, raw_response=True)
+        result = response.json()
+        if task == "audio-to-audio":
+            content = None
+            type = None
+            for k, v in result[0].items():
+                if k == "blob":
+                    content = base64.b64decode(v.encode("utf-8"))
+                if k == "content-type":
+                    type = "audio/flac".split("/")[-1]
+            audio = AudioSegment.from_file(BytesIO(content))
+            name = str(uuid.uuid4())[:4]
+            audio.export(f"public/audios/{name}.{type}", format=type)
+            result = {"generated audio": f"/audios/{name}.{type}"}
+    return result
+
+def local_model_inference(model_id, data, task):
+    inference = partial(models, model_id)
+    # contronlet
+    if model_id.startswith("lllyasviel/sd-controlnet-"):
+        img_url = data["image"]
+        text = data["text"]
+        results = inference({"img_url": img_url, "text": text})
+        if "path" in results:
+            results["generated image"] = results.pop("path")
+        return results
+    if model_id.endswith("-control"):
+        img_url = data["image"]
+        results = inference({"img_url": img_url})
+        if "path" in results:
+            results["generated image"] = results.pop("path")
+        return results
+        
+    if task == "text-to-video":
+        results = inference(data)
+        if "path" in results:
+            results["generated video"] = results.pop("path")
+        return results
+
+    # NLP tasks
+    if task == "question-answering" or task == "sentence-similarity":
+        results = inference(json=data)
+        return results
+    if task in ["text-classification",  "token-classification", "text2text-generation", "summarization", "translation", "conversational", "text-generation"]:
+        results = inference(json=data)
+        return results
+
+    # CV tasks
+    if task == "depth-estimation":
+        img_url = data["image"]
+        results = inference({"img_url": img_url})
+        if "path" in results:
+            results["generated depth image"] = results.pop("path")
+        return results
+    if task == "image-segmentation":
+        img_url = data["image"]
+        results = inference({"img_url": img_url})
+        results["generated image with segmentation mask"] = results.pop("path")
+        return results
+    if task == "image-to-image":
+        img_url = data["image"]
+        results = inference({"img_url": img_url})
+        if "path" in results:
+            results["generated image"] = results.pop("path")
+        return results
+    if task == "text-to-image":
+        results = inference(data)
+        if "path" in results:
+            results["generated image"] = results.pop("path")
+        return results
+    if task == "object-detection":
+        img_url = data["image"]
+        predicted = inference({"img_url": img_url})
+        if "error" in predicted:
+            return predicted
+        image = load_image(img_url)
+        draw = ImageDraw.Draw(image)
+        labels = list(item['label'] for item in predicted)
+        color_map = {}
+        for label in labels:
+            if label not in color_map:
+                color_map[label] = (random.randint(0, 255), random.randint(0, 100), random.randint(0, 255))
+        for label in predicted:
+            box = label["box"]
+            draw.rectangle(((box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])), outline=color_map[label["label"]], width=2)
+            draw.text((box["xmin"]+5, box["ymin"]-15), label["label"], fill=color_map[label["label"]])
+        name = str(uuid.uuid4())[:4]
+        image.save(f"public/images/{name}.jpg")
+        results = {}
+        results["generated image with predicted box"] = f"/images/{name}.jpg"
+        results["predicted"] = predicted
+        return results
+    if task in ["image-classification", "image-to-text", "document-question-answering", "visual-question-answering"]:
+        img_url = data["image"]
+        text = None
+        if "text" in data:
+            text = data["text"]
+        results = inference({"img_url": img_url, "text": text})
+        return results
+    # AUDIO tasks
+    if task == "text-to-speech":
+        results = inference(data)
+        if "path" in results:
+            results["generated audio"] = results.pop("path")
+        return results
+    if task in ["automatic-speech-recognition", "audio-to-audio", "audio-classification"]:
+        audio_url = data["audio"]
+        results = inference({"audio_url": audio_url})
+        return results
+
+
+def model_inference(model_id, data, hosted_on, task, huggingfacetoken=None):
+    if huggingfacetoken:
+        HUGGINGFACE_HEADERS = {
+            "Authorization": f"Bearer {huggingfacetoken}",
+        }
+    else:
+        HUGGINGFACE_HEADERS = None
+    if hosted_on == "unknown":
+        r = status(model_id)
+        logger.debug("Local Server Status: " + str(r))
+        if "loaded" in r and r["loaded"]:
+            hosted_on = "local"
+        else:
+            huggingfaceStatusUrl = f"https://api-inference.huggingface.co/status/{model_id}"
+            r = requests.get(huggingfaceStatusUrl, headers=HUGGINGFACE_HEADERS, proxies=PROXY)
+            logger.debug("Huggingface Status: " + str(r.json()))
+            if "loaded" in r and r["loaded"]:
+                hosted_on = "huggingface"
+    try:
+        if hosted_on == "local":
+            inference_result = local_model_inference(model_id, data, task)
+        elif hosted_on == "huggingface":
+            inference_result = huggingface_model_inference(model_id, data, task, huggingfacetoken)
+    except Exception as e:
+        print(e)
+        traceback.print_exc()
+        inference_result = {"error":{"message": str(e)}}
+    return inference_result
+
+
+def get_model_status(model_id, url, headers, queue = None):
+    endpoint_type = "huggingface" if "huggingface" in url else "local"
+    if "huggingface" in url:
+        r = requests.get(url, headers=headers, proxies=PROXY)
+    else:
+        r = status(model_id)
+    if "loaded" in r and r["loaded"]:
+        if queue:
+            queue.put((model_id, True, endpoint_type))
+        return True
+    else:
+        if queue:
+            queue.put((model_id, False, None))
+        return False
+
+def get_avaliable_models(candidates, topk=10, huggingfacetoken = None):
+    all_available_models = {"local": [], "huggingface": []}
+    threads = []
+    result_queue = Queue()
+    HUGGINGFACE_HEADERS = {
+        "Authorization": f"Bearer {huggingfacetoken}",
+    }
+    for candidate in candidates:
+        model_id = candidate["id"]
+
+        if inference_mode != "local":
+            huggingfaceStatusUrl = f"https://api-inference.huggingface.co/status/{model_id}"
+            thread = threading.Thread(target=get_model_status, args=(model_id, huggingfaceStatusUrl, HUGGINGFACE_HEADERS, result_queue))
+            threads.append(thread)
+            thread.start()
+        
+        if inference_mode != "huggingface" and config["local_deployment"] != "minimal":
+            thread = threading.Thread(target=get_model_status, args=(model_id, "", {}, result_queue))
+            threads.append(thread)
+            thread.start()
+        
+    result_count = len(threads)
+    while result_count:
+        model_id, status, endpoint_type = result_queue.get()
+        if status and model_id not in all_available_models:
+            all_available_models[endpoint_type].append(model_id)
+        if len(all_available_models["local"] + all_available_models["huggingface"]) >= topk:
+            break
+        result_count -= 1
+
+    for thread in threads:
+        thread.join()
+
+    return all_available_models
+
+def collect_result(command, choose, inference_result):
+    result = {"task": command}
+    result["inference result"] = inference_result
+    result["choose model result"] = choose
+    logger.debug(f"inference result: {inference_result}")
+    return result
+
+
+def run_task(input, command, results, openaikey = None, huggingfacetoken = None):
+    id = command["id"]
+    args = command["args"]
+    task = command["task"]
+    deps = command["dep"]
+    if deps[0] != -1:
+        dep_tasks = [results[dep] for dep in deps]
+    else:
+        dep_tasks = []
+    
+    logger.debug(f"Run task: {id} - {task}")
+    logger.debug("Deps: " + json.dumps(dep_tasks))
+
+    if deps[0] != -1:
+        if "image" in args and "<GENERATED>-" in args["image"]:
+            resource_id = int(args["image"].split("-")[1])
+            if "generated image" in results[resource_id]["inference result"]:
+                args["image"] = results[resource_id]["inference result"]["generated image"]
+        if "audio" in args and "<GENERATED>-" in args["audio"]:
+            resource_id = int(args["audio"].split("-")[1])
+            if "generated audio" in results[resource_id]["inference result"]:
+                args["audio"] = results[resource_id]["inference result"]["generated audio"]
+        if "text" in args and "<GENERATED>-" in args["text"]:
+            resource_id = int(args["text"].split("-")[1])
+            if "generated text" in results[resource_id]["inference result"]:
+                args["text"] = results[resource_id]["inference result"]["generated text"]
+
+    text = image = audio = None
+    for dep_task in dep_tasks:
+        if "generated text" in dep_task["inference result"]:
+            text = dep_task["inference result"]["generated text"]
+            logger.debug("Detect the generated text of dependency task (from results):" + text)
+        elif "text" in dep_task["task"]["args"]:
+            text = dep_task["task"]["args"]["text"]
+            logger.debug("Detect the text of dependency task (from args): " + text)
+        if "generated image" in dep_task["inference result"]:
+            image = dep_task["inference result"]["generated image"]
+            logger.debug("Detect the generated image of dependency task (from results): " + image)
+        elif "image" in dep_task["task"]["args"]:
+            image = dep_task["task"]["args"]["image"]
+            logger.debug("Detect the image of dependency task (from args): " + image)
+        if "generated audio" in dep_task["inference result"]:
+            audio = dep_task["inference result"]["generated audio"]
+            logger.debug("Detect the generated audio of dependency task (from results): " + audio)
+        elif "audio" in dep_task["task"]["args"]:
+            audio = dep_task["task"]["args"]["audio"]
+            logger.debug("Detect the audio of dependency task (from args): " + audio)
+
+    if "image" in args and "<GENERATED>" in args["image"]:
+        if image:
+            args["image"] = image
+    if "audio" in args and "<GENERATED>" in args["audio"]:
+        if audio:
+            args["audio"] = audio
+    if "text" in args and "<GENERATED>" in args["text"]:
+        if text:
+            args["text"] = text
+
+    for resource in ["image", "audio"]:
+        if resource in args and not args[resource].startswith("public/") and len(args[resource]) > 0 and not args[resource].startswith("http"):
+            args[resource] = f"public/{args[resource]}"
+    
+    if "-text-to-image" in command['task'] and "text" not in args:
+        logger.debug("control-text-to-image task, but text is empty, so we use control-generation instead.")
+        control = task.split("-")[0]
+        
+        if control == "seg":
+            task = "image-segmentation"
+            command['task'] = task
+        elif control == "depth":
+            task = "depth-estimation"
+            command['task'] = task
+        else:
+            task = f"{control}-control"
+
+    command["args"] = args
+    logger.debug(f"parsed task: {command}")
+
+    if task.endswith("-text-to-image") or task.endswith("-control"):
+        if inference_mode != "huggingface":
+            if task.endswith("-text-to-image"):
+                control = task.split("-")[0]
+                best_model_id = f"lllyasviel/sd-controlnet-{control}"
+            else:
+                best_model_id = task
+            hosted_on = "local"
+            reason = "ControlNet is the best model for this task."
+            choose = {"id": best_model_id, "reason": reason}
+            logger.debug(f"chosen model: {choose}")
+        else:
+            logger.warning(f"Task {command['task']} is not available. ControlNet need to be deployed locally.")
+            record_case(success=False, **{"input": input, "task": command, "reason": f"Task {command['task']} is not available. ControlNet need to be deployed locally.", "op":"message"})
+            inference_result = {"error": f"service related to ControlNet is not available."}
+            results[id] = collect_result(command, "", inference_result)
+            return False
+    elif task in ["summarization", "translation", "conversational", "text-generation", "text2text-generation"]: # ChatGPT Can do
+        best_model_id = "ChatGPT"
+        reason = "ChatGPT performs well on some NLP tasks as well."
+        choose = {"id": best_model_id, "reason": reason}
+        messages = [{
+            "role": "user",
+            "content": f"[ {input} ] contains a task in JSON format {command}, 'task' indicates the task type and 'args' indicates the arguments required for the task. Don't explain the task to me, just help me do it and give me the result. The result must be in text form without any urls."
+        }]
+        response = chitchat(messages, openaikey)
+        results[id] = collect_result(command, choose, {"response": response})
+        return True
+    else:
+        if task not in MODELS_MAP:
+            logger.warning(f"no available models on {task} task.")
+            record_case(success=False, **{"input": input, "task": command, "reason": f"task not support: {command['task']}", "op":"message"})
+            inference_result = {"error": f"{command['task']} not found in available tasks."}
+            results[id] = collect_result(command, "", inference_result)
+            return False
+
+        candidates = MODELS_MAP[task][:20]
+        all_avaliable_models = get_avaliable_models(candidates, config["num_candidate_models"], huggingfacetoken)
+        all_avaliable_model_ids = all_avaliable_models["local"] + all_avaliable_models["huggingface"]
+        logger.debug(f"avaliable models on {command['task']}: {all_avaliable_models}")
+
+        if len(all_avaliable_model_ids) == 0:
+            logger.warning(f"no available models on {command['task']}")
+            record_case(success=False, **{"input": input, "task": command, "reason": f"no available models: {command['task']}", "op":"message"})
+            inference_result = {"error": f"no available models on {command['task']} task."}
+            results[id] = collect_result(command, "", inference_result)
+            return False
+        
+        all_avaliable_model_ids = all_avaliable_model_ids[:1]
+        if len(all_avaliable_model_ids) == 1:
+            best_model_id = all_avaliable_model_ids[0]
+            hosted_on = "local" if best_model_id in all_avaliable_models["local"] else "huggingface"
+            reason = "Only one model available."
+            choose = {"id": best_model_id, "reason": reason}
+            logger.debug(f"chosen model: {choose}")
+        else:
+            cand_models_info = [
+                {
+                    "id": model["id"],
+                    "inference endpoint": all_avaliable_models.get(
+                        "local" if model["id"] in all_avaliable_models["local"] else "huggingface"
+                    ),
+                    "likes": model.get("likes"),
+                    "description": model.get("description", "")[:config["max_description_length"]],
+                    "language": model.get("language"),
+                    "tags": model.get("tags"),
+                }
+                for model in candidates
+                if model["id"] in all_avaliable_model_ids
+            ]
+
+            choose_str = choose_model(input, command, cand_models_info, openaikey)
+            logger.debug(f"chosen model: {choose_str}")
+            try:
+                choose = json.loads(choose_str)
+                reason = choose["reason"]
+                best_model_id = choose["id"]
+                hosted_on = "local" if best_model_id in all_avaliable_models["local"] else "huggingface"
+            except Exception as e:
+                logger.warning(f"the response [ {choose_str} ] is not a valid JSON, try to find the model id and reason in the response.")
+                choose_str = find_json(choose_str)
+                best_model_id, reason, choose  = get_id_reason(choose_str)
+                hosted_on = "local" if best_model_id in all_avaliable_models["local"] else "huggingface"
+    inference_result = model_inference(best_model_id, args, hosted_on, command['task'], huggingfacetoken)
+
+    if "error" in inference_result:
+        logger.warning(f"Inference error: {inference_result['error']}")
+        record_case(success=False, **{"input": input, "task": command, "reason": f"inference error: {inference_result['error']}", "op":"message"})
+        results[id] = collect_result(command, choose, inference_result)
+        return False
+    
+    results[id] = collect_result(command, choose, inference_result)
+    return True
+
+def chat_huggingface(messages, openaikey = None, huggingfacetoken = None, return_planning = False, return_results = False):
+    start = time.time()
+    context = messages[:-1]
+    input = messages[-1]["content"]
+    logger.info("*"*80)
+    logger.info(f"input: {input}")
+
+    task_str = parse_task(context, input, openaikey)
+    logger.info(task_str)
+
+    if "error" in task_str:
+        return str(task_str), {}
+    else:
+        task_str = task_str.strip()
+
+    try:
+        tasks = json.loads(task_str)
+    except Exception as e:
+        logger.debug(e)
+        response = chitchat(messages, openaikey)
+        record_case(success=False, **{"input": input, "task": task_str, "reason": "task parsing fail", "op":"chitchat"})
+        return response, {}
+
+    if task_str == "[]":  # using LLM response for empty task
+        record_case(success=False, **{"input": input, "task": [], "reason": "task parsing fail: empty", "op": "chitchat"})
+        response = chitchat(messages, openaikey)
+        return response, {}
+
+    if len(tasks)==1 and tasks[0]["task"] in ["summarization", "translation", "conversational", "text-generation", "text2text-generation"]:
+        record_case(success=True, **{"input": input, "task": tasks, "reason": "task parsing fail: empty", "op": "chitchat"})
+        response = chitchat(messages, openaikey)
+        best_model_id = "ChatGPT"
+        reason = "ChatGPT performs well on some NLP tasks as well."
+        choose = {"id": best_model_id, "reason": reason}
+        return response, collect_result(tasks[0], choose, {"response": response})
+    
+
+    tasks = unfold(tasks)
+    tasks = fix_dep(tasks)
+    logger.debug(tasks)
+    
+    if return_planning:
+        return tasks
+
+    results = {}
+    threads = []
+    tasks = tasks[:]
+    d = dict()
+    retry = 0
+    while True:
+        num_threads = len(threads)
+        for task in tasks:
+            dep = task["dep"]
+            # logger.debug(f"d.keys(): {d.keys()}, dep: {dep}")
+            for dep_id in dep:
+                if dep_id >= task["id"]:
+                    task["dep"] = [-1]
+                    dep = [-1]
+                    break
+            if len(list(set(dep).intersection(d.keys()))) == len(dep) or dep[0] == -1:
+                tasks.remove(task)
+                thread = threading.Thread(target=run_task, args=(input, task, d, openaikey, huggingfacetoken))
+                thread.start()
+                threads.append(thread)
+        if num_threads == len(threads):
+            time.sleep(0.5)
+            retry += 1
+        if retry > 80:
+            logger.debug("User has waited too long, Loop break.")
+            break
+        if len(tasks) == 0:
+            break
+    for thread in threads:
+        thread.join()
+    
+    results = d.copy()
+
+    logger.debug(results)
+    if return_results:
+        return results
+    
+    response = response_results(input, results, openaikey).strip()
+
+    end = time.time()
+    during = end - start
+
+    answer = {"message": response}
+    record_case(success=True, **{"input": input, "task": task_str, "results": results, "response": response, "during": during, "op":"response"})
+    logger.info(f"response: {response}")
+    return response, results