modity code structure

src/llmtuner/api/app.py

@@ -0,0 +1,152 @@
+import uvicorn
+from threading import Thread
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from transformers import TextIteratorStreamer
+from contextlib import asynccontextmanager
+from sse_starlette import EventSourceResponse
+from typing import Any, Dict
+
+from llmtuner.tuner import get_infer_args, load_model_and_tokenizer
+from llmtuner.extras.misc import get_logits_processor, torch_gc
+from llmtuner.extras.template import Template
+from llmtuner.api.protocol import (
+    ModelCard,
+    ModelList,
+    ChatMessage,
+    DeltaMessage,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionStreamResponse,
+    ChatCompletionResponseChoice,
+    ChatCompletionResponseStreamChoice,
+    ChatCompletionResponseUsage
+)
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI): # collects GPU memory
+    yield
+    torch_gc()
+
+
+def create_app():
+    model_args, data_args, finetuning_args, generating_args = get_infer_args()
+    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
+
+    prompt_template = Template(data_args.prompt_template)
+    source_prefix = data_args.source_prefix if data_args.source_prefix else ""
+
+    app = FastAPI(lifespan=lifespan)
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    @app.get("/v1/models", response_model=ModelList)
+    async def list_models():
+        global model_args
+        model_card = ModelCard(id="gpt-3.5-turbo")
+        return ModelList(data=[model_card])
+
+    @app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
+    async def create_chat_completion(request: ChatCompletionRequest):
+        if request.messages[-1].role != "user":
+            raise HTTPException(status_code=400, detail="Invalid request")
+        query = request.messages[-1].content
+
+        prev_messages = request.messages[:-1]
+        if len(prev_messages) > 0 and prev_messages[0].role == "system":
+            prefix = prev_messages.pop(0).content
+        else:
+            prefix = source_prefix
+
+        history = []
+        if len(prev_messages) % 2 == 0:
+            for i in range(0, len(prev_messages), 2):
+                if prev_messages[i].role == "user" and prev_messages[i+1].role == "assistant":
+                    history.append([prev_messages[i].content, prev_messages[i+1].content])
+
+        inputs = tokenizer([prompt_template.get_prompt(query, history, prefix)], return_tensors="pt")
+        inputs = inputs.to(model.device)
+
+        gen_kwargs = generating_args.to_dict()
+        gen_kwargs.update({
+            "input_ids": inputs["input_ids"],
+            "temperature": request.temperature if request.temperature else gen_kwargs["temperature"],
+            "top_p": request.top_p if request.top_p else gen_kwargs["top_p"],
+            "logits_processor": get_logits_processor()
+        })
+
+        if request.max_tokens:
+            gen_kwargs.pop("max_length", None)
+            gen_kwargs["max_new_tokens"] = request.max_tokens
+
+        if request.stream:
+            generate = predict(gen_kwargs, request.model)
+            return EventSourceResponse(generate, media_type="text/event-stream")
+
+        generation_output = model.generate(**gen_kwargs)
+        outputs = generation_output.tolist()[0][len(inputs["input_ids"][0]):]
+        response = tokenizer.decode(outputs, skip_special_tokens=True)
+
+        usage = ChatCompletionResponseUsage(
+            prompt_tokens=len(inputs["input_ids"][0]),
+            completion_tokens=len(outputs),
+            total_tokens=len(inputs["input_ids"][0]) + len(outputs)
+        )
+
+        choice_data = ChatCompletionResponseChoice(
+            index=0,
+            message=ChatMessage(role="assistant", content=response),
+            finish_reason="stop"
+        )
+
+        return ChatCompletionResponse(model=request.model, choices=[choice_data], usage=usage, object="chat.completion")
+
+    async def predict(gen_kwargs: Dict[str, Any], model_id: str):
+        streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+        gen_kwargs["streamer"] = streamer
+
+        thread = Thread(target=model.generate, kwargs=gen_kwargs)
+        thread.start()
+
+        choice_data = ChatCompletionResponseStreamChoice(
+            index=0,
+            delta=DeltaMessage(role="assistant"),
+            finish_reason=None
+        )
+        chunk = ChatCompletionStreamResponse(model=model_id, choices=[choice_data], object="chat.completion.chunk")
+        yield chunk.json(exclude_unset=True, ensure_ascii=False)
+
+        for new_text in streamer:
+            if len(new_text) == 0:
+                continue
+
+            choice_data = ChatCompletionResponseStreamChoice(
+                index=0,
+                delta=DeltaMessage(content=new_text),
+                finish_reason=None
+            )
+            chunk = ChatCompletionStreamResponse(model=model_id, choices=[choice_data], object="chat.completion.chunk")
+            yield chunk.json(exclude_unset=True, ensure_ascii=False)
+
+        choice_data = ChatCompletionResponseStreamChoice(
+            index=0,
+            delta=DeltaMessage(),
+            finish_reason="stop"
+        )
+        chunk = ChatCompletionStreamResponse(model=model_id, choices=[choice_data], object="chat.completion.chunk")
+        yield chunk.json(exclude_unset=True, ensure_ascii=False)
+        yield "[DONE]"
+
+    return app
+
+
+if __name__ == "__main__":
+    app = create_app()
+    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)