wehub-resource-sync
138 lines
4.1 KiB
Python
138 lines
4.1 KiB
Python
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import random
|
|
import string
|
|
import time
|
|
|
|
import paddle
|
|
import uvicorn
|
|
from fastapi import FastAPI, Response, status
|
|
from pydantic import BaseModel
|
|
from sse_starlette.sse import EventSourceResponse
|
|
|
|
from paddlenlp.transformers import CodeGenForCausalLM, CodeGenTokenizer
|
|
from paddlenlp.utils.log import logger
|
|
|
|
|
|
class DefaultConfig:
|
|
model_name_or_path = "Salesforce/codegen-350M-mono"
|
|
device = "gpu"
|
|
temperature = 0.5
|
|
top_k = 10
|
|
top_p = 1.0
|
|
repetition_penalty = 1.0
|
|
min_length = 0
|
|
max_length = 16
|
|
decode_strategy = "greedy_search"
|
|
use_faster = True
|
|
use_fp16_decoding = True
|
|
default_dtype = "float16" if use_faster and use_fp16_decoding else "float32"
|
|
|
|
|
|
class Input(BaseModel):
|
|
prompt: str
|
|
stream: bool = False
|
|
|
|
|
|
class Output(BaseModel):
|
|
id: str
|
|
model: str = "codegen"
|
|
object: str = "text_completion"
|
|
created: int = int(time.time())
|
|
choices: list = None
|
|
usage = {
|
|
"completion_tokens": None,
|
|
"prompt_tokens": None,
|
|
"total_tokens": None,
|
|
}
|
|
|
|
|
|
generate_config = DefaultConfig()
|
|
paddle.set_device(generate_config.device)
|
|
paddle.set_default_dtype(generate_config.default_dtype)
|
|
|
|
tokenizer = CodeGenTokenizer.from_pretrained(generate_config.model_name_or_path)
|
|
model = CodeGenForCausalLM.from_pretrained(generate_config.model_name_or_path)
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
def random_completion_id():
|
|
return "cmpl-" + "".join(random.choice(string.ascii_letters + string.digits) for _ in range(29))
|
|
|
|
|
|
@app.post("/v1/engines/codegen/completions", status_code=200)
|
|
async def gen(item: Input):
|
|
item = item.dict()
|
|
logger.info(f"Request: {item}")
|
|
temperature = item.get("temperature", generate_config.temperature)
|
|
top_k = item.get("top_k", generate_config.top_k)
|
|
if temperature == 0.0:
|
|
temperature = 1.0
|
|
top_k = 1
|
|
repetition_penalty = item.get("frequency_penalty", generate_config.repetition_penalty)
|
|
|
|
start_time = time.time()
|
|
logger.info("Start generating code")
|
|
tokenized = tokenizer([item["prompt"]], truncation=True, return_tensors="pd")
|
|
output, _ = model.generate(
|
|
tokenized["input_ids"],
|
|
max_length=16,
|
|
min_length=generate_config.min_length,
|
|
decode_strategy=generate_config.decode_strategy,
|
|
top_k=top_k,
|
|
repetition_penalty=repetition_penalty,
|
|
temperature=temperature,
|
|
use_fast=generate_config.use_faster,
|
|
use_fp16_decoding=generate_config.use_fp16_decoding,
|
|
)
|
|
logger.info("Finish generating code")
|
|
end_time = time.time()
|
|
logger.info(f"Time cost: {end_time - start_time}")
|
|
output = tokenizer.decode(output[0], skip_special_tokens=True)
|
|
logger.info(f"Generated code: {output}")
|
|
output_json = Output(
|
|
id=random_completion_id(),
|
|
choices=[
|
|
{
|
|
"text": output,
|
|
"index": 0,
|
|
"finish_reason": "stop",
|
|
"logprobs": None,
|
|
}
|
|
],
|
|
usage={
|
|
"completion_tokens": None,
|
|
"prompt_tokens": None,
|
|
"total_tokens": None,
|
|
},
|
|
).json()
|
|
|
|
def stream_response(response):
|
|
yield f"{response}\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
|
|
if item.get("stream", False):
|
|
return EventSourceResponse(stream_response(output_json))
|
|
else:
|
|
return Response(
|
|
status_code=status.HTTP_200_OK,
|
|
content=output_json,
|
|
media_type="application/json",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run("codegen_server:app", host="0.0.0.0", port=8978)
|