llmbb/LLMBB-Agent

1# coding=utf-82# Implements API for Qwen-7B in OpenAI's format. (https://platform.openai.com/docs/api-reference/chat)3# Usage: python openai_api.py4# Visit http://localhost:8000/docs for documents.5 6import re7import copy8import json9import time10from argparse import ArgumentParser11from contextlib import asynccontextmanager12from typing import Dict, List, Literal, Optional, Union13 14import torch15import uvicorn16from fastapi import FastAPI, HTTPException17from fastapi.middleware.cors import CORSMiddleware18from pydantic import BaseModel, Field19from sse_starlette.sse import EventSourceResponse20from transformers import AutoTokenizer, AutoModelForCausalLM21from transformers.generation import GenerationConfig22from starlette.middleware.base import BaseHTTPMiddleware23from starlette.requests import Request24from starlette.responses import Response25import base6426 27 28class BasicAuthMiddleware(BaseHTTPMiddleware):29    def __init__(self, app, username: str, password: str):30        super().__init__(app)31        self.required_credentials = base64.b64encode(f"{username}:{password}".encode()).decode()32 33    async def dispatch(self, request: Request, call_next):34        authorization: str = request.headers.get("Authorization")35        if authorization:36            try:37                schema, credentials = authorization.split()38                if credentials == self.required_credentials:39                    return await call_next(request)40            except ValueError:41                pass42 43        headers = {'WWW-Authenticate': 'Basic'}44        return Response(status_code=401, headers=headers)45 46 47def _gc(forced: bool = False):48    global args49    if args.disable_gc and not forced:50        return51 52    import gc53    gc.collect()54    if torch.cuda.is_available():55        torch.cuda.empty_cache()56 57 58@asynccontextmanager59async def lifespan(app: FastAPI):  # collects GPU memory60    yield61    _gc(forced=True)62 63 64app = FastAPI(lifespan=lifespan)65 66app.add_middleware(67    CORSMiddleware,68    allow_origins=["*"],69    allow_credentials=True,70    allow_methods=["*"],71    allow_headers=["*"],72)73 74 75class ModelCard(BaseModel):76    id: str77    object: str = "model"78    created: int = Field(default_factory=lambda: int(time.time()))79    owned_by: str = "owner"80    root: Optional[str] = None81    parent: Optional[str] = None82    permission: Optional[list] = None83 84 85class ModelList(BaseModel):86    object: str = "list"87    data: List[ModelCard] = []88 89 90class ChatMessage(BaseModel):91    role: Literal["user", "assistant", "system", "function"]92    content: Optional[str]93    function_call: Optional[Dict] = None94 95 96class DeltaMessage(BaseModel):97    role: Optional[Literal["user", "assistant", "system"]] = None98    content: Optional[str] = None99 100 101class ChatCompletionRequest(BaseModel):102    model: str103    messages: List[ChatMessage]104    functions: Optional[List[Dict]] = None105    temperature: Optional[float] = None106    top_p: Optional[float] = None107    max_length: Optional[int] = None108    stream: Optional[bool] = False109    stop: Optional[List[str]] = None110 111 112class ChatCompletionResponseChoice(BaseModel):113    index: int114    message: ChatMessage115    finish_reason: Literal["stop", "length", "function_call"]116 117 118class ChatCompletionResponseStreamChoice(BaseModel):119    index: int120    delta: DeltaMessage121    finish_reason: Optional[Literal["stop", "length"]]122 123 124class ChatCompletionResponse(BaseModel):125    model: str126    object: Literal["chat.completion", "chat.completion.chunk"]127    choices: List[128        Union[ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice]129    ]130    created: Optional[int] = Field(default_factory=lambda: int(time.time()))131 132 133@app.get("/v1/models", response_model=ModelList)134async def list_models():135    global model_args136    model_card = ModelCard(id="gpt-3.5-turbo")137    return ModelList(data=[model_card])138 139 140# To work around that unpleasant leading-\n tokenization issue!141def add_extra_stop_words(stop_words):142    if stop_words:143        _stop_words = []144        _stop_words.extend(stop_words)145        for x in stop_words:146            s = x.lstrip("\n")147            if s and (s not in _stop_words):148                _stop_words.append(s)149        return _stop_words150    return stop_words151 152 153def trim_stop_words(response, stop_words):154    if stop_words:155        for stop in stop_words:156            idx = response.find(stop)157            if idx != -1:158                response = response[:idx]159    return response160 161 162TOOL_DESC = """{name_for_model}: Call this tool to interact with the {name_for_human} API. What is the {name_for_human} API useful for? {description_for_model} Parameters: {parameters}"""163 164REACT_INSTRUCTION = """Answer the following questions as best you can. You have access to the following APIs:165 166{tools_text}167 168Use the following format:169 170Question: the input question you must answer171Thought: you should always think about what to do172Action: the action to take, should be one of [{tools_name_text}]173Action Input: the input to the action174Observation: the result of the action175... (this Thought/Action/Action Input/Observation can be repeated zero or more times)176Thought: I now know the final answer177Final Answer: the final answer to the original input question178 179Begin!"""180 181_TEXT_COMPLETION_CMD = object()182 183 184#185# Temporarily, the system role does not work as expected.186# We advise that you write the setups for role-play in your query,187# i.e., use the user role instead of the system role.188#189# TODO: Use real system role when the model is ready.190#191def parse_messages(messages, functions):192    if all(m.role != "user" for m in messages):193        raise HTTPException(194            status_code=400,195            detail=f"Invalid request: Expecting at least one user message.",196        )197 198    messages = copy.deepcopy(messages)199    default_system = "You are a helpful assistant."200    system = ""201    if messages[0].role == "system":202        system = messages.pop(0).content.lstrip("\n").rstrip()203        if system == default_system:204            system = ""205 206    if functions:207        tools_text = []208        tools_name_text = []209        for func_info in functions:210            name = func_info.get("name", "")211            name_m = func_info.get("name_for_model", name)212            name_h = func_info.get("name_for_human", name)213            desc = func_info.get("description", "")214            desc_m = func_info.get("description_for_model", desc)215            tool = TOOL_DESC.format(216                name_for_model=name_m,217                name_for_human=name_h,218                # Hint: You can add the following format requirements in description:219                #   "Format the arguments as a JSON object."220                #   "Enclose the code within triple backticks (`) at the beginning and end of the code."221                description_for_model=desc_m,222                parameters=json.dumps(func_info["parameters"], ensure_ascii=False),223            )224            tools_text.append(tool)225            tools_name_text.append(name_m)226        tools_text = "\n\n".join(tools_text)227        tools_name_text = ", ".join(tools_name_text)228        system += "\n\n" + REACT_INSTRUCTION.format(229            tools_text=tools_text,230            tools_name_text=tools_name_text,231        )232        system = system.lstrip("\n").rstrip()233 234    dummy_thought = {235        "en": "\nThought: I now know the final answer.\nFinal answer: ",236        "zh": "\nThought: 我会作答了。\nFinal answer: ",237    }238 239    _messages = messages240    messages = []241    for m_idx, m in enumerate(_messages):242        role, content, func_call = m.role, m.content, m.function_call243        if content:244            content = content.lstrip("\n").rstrip()245        if role == "function":246            if (len(messages) == 0) or (messages[-1].role != "assistant"):247                raise HTTPException(248                    status_code=400,249                    detail=f"Invalid request: Expecting role assistant before role function.",250                )251            messages[-1].content += f"\nObservation: {content}"252            if m_idx == len(_messages) - 1:253                messages[-1].content += "\nThought:"254        elif role == "assistant":255            if len(messages) == 0:256                raise HTTPException(257                    status_code=400,258                    detail=f"Invalid request: Expecting role user before role assistant.",259                )260            last_msg = messages[-1].content261            last_msg_has_zh = len(re.findall(r"[\u4e00-\u9fff]+", last_msg)) > 0262            if func_call is None:263                if functions:264                    content = dummy_thought["zh" if last_msg_has_zh else "en"] + content265            else:266                f_name, f_args = func_call["name"], func_call["arguments"]267                if not content:268                    if last_msg_has_zh:269                        content = f"Thought: 我可以使用 {f_name} API。"270                    else:271                        content = f"Thought: I can use {f_name}."272                content = f"\n{content}\nAction: {f_name}\nAction Input: {f_args}"273            if messages[-1].role == "user":274                messages.append(275                    ChatMessage(role="assistant", content=content.lstrip("\n").rstrip())276                )277            else:278                messages[-1].content += content279        elif role == "user":280            messages.append(281                ChatMessage(role="user", content=content.lstrip("\n").rstrip())282            )283        else:284            raise HTTPException(285                status_code=400, detail=f"Invalid request: Incorrect role {role}."286            )287 288    query = _TEXT_COMPLETION_CMD289    if messages[-1].role == "user":290        query = messages[-1].content291        messages = messages[:-1]292 293    if len(messages) % 2 != 0:294        raise HTTPException(status_code=400, detail="Invalid request")295 296    history = []  # [(Q1, A1), (Q2, A2), ..., (Q_last_turn, A_last_turn)]297    for i in range(0, len(messages), 2):298        if messages[i].role == "user" and messages[i + 1].role == "assistant":299            usr_msg = messages[i].content.lstrip("\n").rstrip()300            bot_msg = messages[i + 1].content.lstrip("\n").rstrip()301            if system and (i == len(messages) - 2):302                usr_msg = f"{system}\n\nQuestion: {usr_msg}"303                system = ""304            for t in dummy_thought.values():305                t = t.lstrip("\n")306                if bot_msg.startswith(t) and ("\nAction: " in bot_msg):307                    bot_msg = bot_msg[len(t):]308            history.append([usr_msg, bot_msg])309        else:310            raise HTTPException(311                status_code=400,312                detail="Invalid request: Expecting exactly one user (or function) role before every assistant role.",313            )314    if system:315        assert query is not _TEXT_COMPLETION_CMD316        query = f"{system}\n\nQuestion: {query}"317    return query, history318 319 320def parse_response(response):321    func_name, func_args = "", ""322    i = response.rfind("\nAction:")323    j = response.rfind("\nAction Input:")324    k = response.rfind("\nObservation:")325    if 0 <= i < j:  # If the text has `Action` and `Action input`,326        if k < j:  # but does not contain `Observation`,327            # then it is likely that `Observation` is omitted by the LLM,328            # because the output text may have discarded the stop word.329            response = response.rstrip() + "\nObservation:"  # Add it back.330        k = response.rfind("\nObservation:")331        func_name = response[i + len("\nAction:"): j].strip()332        func_args = response[j + len("\nAction Input:"): k].strip()333    if func_name:334        choice_data = ChatCompletionResponseChoice(335            index=0,336            message=ChatMessage(337                role="assistant",338                content=response[:i],339                function_call={"name": func_name, "arguments": func_args},340            ),341            finish_reason="function_call",342        )343        return choice_data344    z = response.rfind("\nFinal Answer: ")345    if z >= 0:346        response = response[z + len("\nFinal Answer: "):]347    choice_data = ChatCompletionResponseChoice(348        index=0,349        message=ChatMessage(role="assistant", content=response),350        finish_reason="stop",351    )352    return choice_data353 354 355# completion mode, not chat mode356def text_complete_last_message(history, stop_words_ids, gen_kwargs):357    im_start = "<|im_start|>"358    im_end = "<|im_end|>"359    prompt = f"{im_start}system\nYou are a helpful assistant.{im_end}"360    for i, (query, response) in enumerate(history):361        query = query.lstrip("\n").rstrip()362        response = response.lstrip("\n").rstrip()363        prompt += f"\n{im_start}user\n{query}{im_end}"364        prompt += f"\n{im_start}assistant\n{response}{im_end}"365    prompt = prompt[: -len(im_end)]366 367    _stop_words_ids = [tokenizer.encode(im_end)]368    if stop_words_ids:369        for s in stop_words_ids:370            _stop_words_ids.append(s)371    stop_words_ids = _stop_words_ids372 373    input_ids = torch.tensor([tokenizer.encode(prompt)]).to(model.device)374    output = model.generate(input_ids, stop_words_ids=stop_words_ids, **gen_kwargs).tolist()[0]375    output = tokenizer.decode(output, errors="ignore")376    assert output.startswith(prompt)377    output = output[len(prompt):]378    output = trim_stop_words(output, ["<|endoftext|>", im_end])379    print(f"<completion>\n{prompt}\n<!-- *** -->\n{output}\n</completion>")380    return output381 382 383@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)384async def create_chat_completion(request: ChatCompletionRequest):385    global model, tokenizer386 387    gen_kwargs = {}388    if request.temperature is not None:389        if request.temperature < 0.01:390            gen_kwargs['top_k'] = 1  # greedy decoding391        else:392            # Not recommended. Please tune top_p instead.393            gen_kwargs['temperature'] = request.temperature394    if request.top_p is not None:395        gen_kwargs['top_p'] = request.top_p396 397    stop_words = add_extra_stop_words(request.stop)398    if request.functions:399        stop_words = stop_words or []400        if "Observation:" not in stop_words:401            stop_words.append("Observation:")402 403    query, history = parse_messages(request.messages, request.functions)404 405    if request.stream:406        if request.functions:407            raise HTTPException(408                status_code=400,409                detail="Invalid request: Function calling is not yet implemented for stream mode.",410            )411        generate = predict(query, history, request.model, stop_words, gen_kwargs)412        return generate413        # return EventSourceResponse(generate, media_type="text/event-stream")414 415    stop_words_ids = [tokenizer.encode(s) for s in stop_words] if stop_words else None416    if query is _TEXT_COMPLETION_CMD:417        response = text_complete_last_message(history, stop_words_ids=stop_words_ids, gen_kwargs=gen_kwargs)418    else:419        response, _ = model.chat(420            tokenizer,421            query,422            history=history,423            stop_words_ids=stop_words_ids,424            **gen_kwargs425        )426        print(f"<chat>\n{history}\n{query}\n<!-- *** -->\n{response}\n</chat>")427    _gc()428 429    response = trim_stop_words(response, stop_words)430    if request.functions:431        choice_data = parse_response(response)432    else:433        choice_data = ChatCompletionResponseChoice(434            index=0,435            message=ChatMessage(role="assistant", content=response),436            finish_reason="stop",437        )438    return ChatCompletionResponse(439        model=request.model, choices=[choice_data], object="chat.completion"440    )441 442 443def _dump_json(data: BaseModel, *args, **kwargs) -> str:444    try:445        return data.model_dump_json(*args, **kwargs)446    except AttributeError:  # pydantic<2.0.0447        return data.json(*args, **kwargs)  # noqa448 449 450async def predict(451        query: str, history: List[List[str]], model_id: str, stop_words: List[str], gen_kwargs: Dict,452):453    global model, tokenizer454    choice_data = ChatCompletionResponseStreamChoice(455        index=0, delta=DeltaMessage(role="assistant"), finish_reason=None456    )457    chunk = ChatCompletionResponse(458        model=model_id, choices=[choice_data], object="chat.completion.chunk"459    )460    yield "{}".format(_dump_json(chunk, exclude_unset=True))461 462    current_length = 0463    stop_words_ids = [tokenizer.encode(s) for s in stop_words] if stop_words else None464    if stop_words:465        # TODO: It's a little bit tricky to trim stop words in the stream mode.466        raise HTTPException(467            status_code=400,468            detail="Invalid request: custom stop words are not yet supported for stream mode.",469        )470    response_generator = model.chat_stream(471        tokenizer, query, history=history, stop_words_ids=stop_words_ids, **gen_kwargs472    )473    for new_response in response_generator:474        if len(new_response) == current_length:475            continue476 477        new_text = new_response[current_length:]478        current_length = len(new_response)479 480        choice_data = ChatCompletionResponseStreamChoice(481            index=0, delta=DeltaMessage(content=new_text), finish_reason=None482        )483        chunk = ChatCompletionResponse(484            model=model_id, choices=[choice_data], object="chat.completion.chunk"485        )486        yield "{}".format(_dump_json(chunk, exclude_unset=True))487 488    choice_data = ChatCompletionResponseStreamChoice(489        index=0, delta=DeltaMessage(), finish_reason="stop"490    )491    chunk = ChatCompletionResponse(492        model=model_id, choices=[choice_data], object="chat.completion.chunk"493    )494    yield "{}".format(_dump_json(chunk, exclude_unset=True))495    yield "[DONE]"496 497    _gc()498 499 500def _get_args():501    parser = ArgumentParser()502    parser.add_argument(503        "-c",504        "--checkpoint-path",505        type=str,506        default="Qwen/Qwen-7B-Chat",507        help="Checkpoint name or path, default to %(default)r",508    )509    parser.add_argument(510        "--api-auth", help="API authentication credentials"511    )512    parser.add_argument(513        "--cpu-only", action="store_true", help="Run demo with CPU only"514    )515    parser.add_argument(516        "--server-port", type=int, default=8000, help="Demo server port."517    )518    parser.add_argument(519        "--server-name",520        type=str,521        default="127.0.0.1",522        help="Demo server name. Default: 127.0.0.1, which is only visible from the local computer."523             " If you want other computers to access your server, use 0.0.0.0 instead.",524    )525    parser.add_argument("--disable-gc", action="store_true",526                        help="Disable GC after each response generated.")527 528    args = parser.parse_args()529    return args530 531 532if __name__ == "__main__":533    args = _get_args()534 535    tokenizer = AutoTokenizer.from_pretrained(536        args.checkpoint_path,537        trust_remote_code=True,538        resume_download=True,539    )540 541    if args.api_auth:542        app.add_middleware(543            BasicAuthMiddleware, username=args.api_auth.split(":")[0], password=args.api_auth.split(":")[1]544        )545 546    if args.cpu_only:547        device_map = "cpu"548    else:549        device_map = "auto"550 551    model = AutoModelForCausalLM.from_pretrained(552        args.checkpoint_path,553        device_map=device_map,554        trust_remote_code=True,555        resume_download=True,556    ).eval()557 558    model.generation_config = GenerationConfig.from_pretrained(559        args.checkpoint_path,560        trust_remote_code=True,561        resume_download=True,562    )563 564    uvicorn.run(app, host=args.server_name, port=args.server_port, workers=1)