Repository: lm-sys/FastChat
File: fastchat/serve/huggingface_api_worker.py (Python)

Code before:
    return {"context_length": worker.context_len}


def create_model_worker():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="localhost")
    parser.add_argument("--port", type=int, default=21002)
    parser.add_argument("--worker-address", type=str, default="http://localhost:21002")
    parser.add_argument(
        "--controller-address", type=str, default="http://localhost:21001"
    )
    parser.add_argument(
        "--supported-models-file", type=str, default="supported_models.json"
    )
    parser.add_argument(
        "--model",
        type=str,
        default="falcon-180b-chat",
        help="The model name to be called.",
    )
    parser.add_argument(
        "--model-names",
        type=lambda s: s.split(","),
        default="falcon-180b-chat",
        help="Optional display comma separated names",
    )
    parser.add_argument(
        "--conv-template", type=str, default=None, help="Conversation prompt template."
    )
    parser.add_argument("--embed-in-truncate", action="store_true")
    parser.add_argument(
        "--limit-worker-concurrency",
        type=int,
        default=5,
        help="Limit the model concurrency to prevent OOM.",
    )
    parser.add_argument("--stream-interval", type=int, default=2)
    parser.add_argument("--no-register", action="store_true")
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Overwrite the random seed for each generation.",
    )
    args = parser.parse_args()
    logger.info(f"args: {args}")

    with open(args.supported_models_file, "r") as f:
        supported_models = json.load(f)

    if args.model not in supported_models:

Review comment (suggestion):
1. Read all of these from `supported_models.json`.
2. You should be able to support multiple models listed in `supported_models.json`.

Code after:

    model_path_list = []
    api_base_list = []
    token_list = []
    context_length_list = []
    model_names_list = []
    conv_template_list = []

    for m in model_info:
        model_path_list.append(model_info[m]["model_path"])
        api_base_list.append(model_info[m]["api_base"])
        token_list.append(model_info[m]["token"])

        context_length = model_info[m].get("context_length", 1024)
        model_names = model_info[m].get("model_names", [m.split("/")[-1]])
        if isinstance(model_names, str):
            model_names = [model_names]
        conv_template = model_info[m].get("conv_template", None)

        context_length_list.append(context_length)
        model_names_list.append(model_names)
        conv_template_list.append(conv_template)

    logger.info(f"Model paths: {model_path_list}")
    logger.info(f"API bases: {api_base_list}")
    logger.info(f"Tokens: {token_list}")
    logger.info(f"Context lengths: {context_length_list}")
    logger.info(f"Model names: {model_names_list}")
    logger.info(f"Conv templates: {conv_template_list}")

    for (
        model_names,
        conv_template,
        model_path,
        api_base,
        token,
        context_length,
    ) in zip(
        model_names_list,
        conv_template_list,
        model_path_list,
        api_base_list,
        token_list,
        context_length_list,
    ):
        m = HuggingfaceApiWorker(
            args.controller_address,
            args.worker_address,
            worker_id,
            model_path,
            api_base,