From 24dada9d5aeaf7db100168300454f363e7208cf1 Mon Sep 17 00:00:00 2001 From: Whitroom <1062015905@qq.com> Date: Wed, 12 Apr 2023 14:45:52 +0800 Subject: [PATCH 1/8] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E8=AF=B7=E6=B1=82?= =?UTF-8?q?=E4=B8=8E=E5=93=8D=E5=BA=94=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/api.py b/api.py index 693c70a..efc28bc 100644 --- a/api.py +++ b/api.py @@ -1,7 +1,11 @@ -from fastapi import FastAPI, Request -from transformers import AutoTokenizer, AutoModel -import uvicorn, json, datetime +import datetime +import json + import torch +import uvicorn +from fastapi import FastAPI, Request +from pydantic import BaseModel +from transformers import AutoModel, AutoTokenizer DEVICE = "cuda" DEVICE_ID = "0" @@ -17,6 +21,18 @@ def torch_gc(): app = FastAPI() +class Item(BaseModel): + prompt: str + history: list[tuple[str, str]] = [[]] + max_length: int = 2048 + top_p: float = 0.7 + temperature: float = 0.95 + +class Answer(BaseModel): + response: str + history: list[tuple[str, str]] + status: int + time: str @app.post("/") async def create_item(request: Request): From be7e14ce45c6075c5cc6b0d3ffa89bc5777d5162 Mon Sep 17 00:00:00 2001 From: Whitroom <1062015905@qq.com> Date: Wed, 12 Apr 2023 14:48:49 +0800 Subject: [PATCH 2/8] =?UTF-8?q?=E4=BD=BF=E7=94=A8=E5=BC=82=E6=AD=A5?= =?UTF-8?q?=E7=9A=84=E4=B8=8A=E4=B8=8B=E6=96=87=E7=AE=A1=E7=90=86=E5=99=A8?= =?UTF-8?q?=EF=BC=8C=E5=90=AF=E5=8A=A8=E5=8A=A0=E9=80=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api.py | 62 ++++++++++++++++++++++++++-------------------------------- 1 file changed, 28 insertions(+), 34 deletions(-) diff --git a/api.py b/api.py index efc28bc..3e0e533 100644 --- a/api.py +++ b/api.py @@ -1,9 +1,9 @@ import datetime -import json +from contextlib import asynccontextmanager import torch import uvicorn -from fastapi import FastAPI, Request +from fastapi import FastAPI from pydantic import BaseModel from transformers import AutoModel, AutoTokenizer @@ -11,6 +11,21 @@ DEVICE = "cuda" DEVICE_ID = "0" CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE +models = {} + +@asynccontextmanager +async def lifespan(app: FastAPI): + models['chat'] = AutoModel.from_pretrained( + "THUDM/models", + trust_remote_code=True).half().cuda() + models['chat'].eval() + models['tokenizer'] = AutoTokenizer.from_pretrained( + "THUDM/models", + trust_remote_code=True) + yield + for model in models.values(): + del model + torch_gc() def torch_gc(): if torch.cuda.is_available(): @@ -18,8 +33,7 @@ def torch_gc(): torch.cuda.empty_cache() torch.cuda.ipc_collect() - -app = FastAPI() +app = FastAPI(lifespan=lifespan) class Item(BaseModel): prompt: str @@ -35,38 +49,18 @@ class Answer(BaseModel): time: str @app.post("/") -async def create_item(request: Request): - global model, tokenizer - json_post_raw = await request.json() - json_post = json.dumps(json_post_raw) - json_post_list = json.loads(json_post) - prompt = json_post_list.get('prompt') - history = json_post_list.get('history') - max_length = json_post_list.get('max_length') - top_p = json_post_list.get('top_p') - temperature = json_post_list.get('temperature') - response, history = model.chat(tokenizer, - prompt, - history=history, - max_length=max_length if max_length else 2048, - top_p=top_p if top_p else 0.7, - temperature=temperature if temperature else 0.95) +async def create_item(item: Item): + response, history = models['chat'].chat( + models['tokenizer'], + item.prompt, + history=item.history, + max_length=item.max_length, + top_p=item.top_p, + temperature=item.temperature) now = datetime.datetime.now() time = now.strftime("%Y-%m-%d %H:%M:%S") - answer = { - "response": response, - "history": history, - "status": 200, - "time": time - } - log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"' - print(log) - torch_gc() - return answer - + print(f"[{time}] prompt: '{item.prompt}', response: '{response}'") + return Answer(response=response, history=history, status=200, time=time) if __name__ == '__main__': - tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True) - model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda() - model.eval() uvicorn.run(app, host='0.0.0.0', port=8000, workers=1) From 02d4d3dd2c660fc90432aa72719e1a7e6e483e41 Mon Sep 17 00:00:00 2001 From: Whitroom <1062015905@qq.com> Date: Wed, 12 Apr 2023 15:03:33 +0800 Subject: [PATCH 3/8] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E6=A8=A1=E5=9E=8B?= =?UTF-8?q?=E4=BD=8D=E7=BD=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/api.py b/api.py index 3e0e533..8997712 100644 --- a/api.py +++ b/api.py @@ -16,11 +16,11 @@ models = {} @asynccontextmanager async def lifespan(app: FastAPI): models['chat'] = AutoModel.from_pretrained( - "THUDM/models", + "THUDM/chatglm-6b", trust_remote_code=True).half().cuda() models['chat'].eval() models['tokenizer'] = AutoTokenizer.from_pretrained( - "THUDM/models", + "THUDM/chatglm-6b", trust_remote_code=True) yield for model in models.values(): @@ -63,4 +63,4 @@ async def create_item(item: Item): return Answer(response=response, history=history, status=200, time=time) if __name__ == '__main__': - uvicorn.run(app, host='0.0.0.0', port=8000, workers=1) + uvicorn.run(app, host='0.0.0.0', port=8010, workers=1) From e63afa6f4503b06edeabee939d8189b202785c7a Mon Sep 17 00:00:00 2001 From: Whitroom <1062015905@qq.com> Date: Wed, 12 Apr 2023 15:28:24 +0800 Subject: [PATCH 4/8] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=8E=A5=E5=8F=A3?= =?UTF-8?q?=E6=96=87=E6=A1=A3=E7=9A=84=E6=B3=A8=E9=87=8A?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api.py b/api.py index 8997712..d706109 100644 --- a/api.py +++ b/api.py @@ -36,7 +36,7 @@ def torch_gc(): app = FastAPI(lifespan=lifespan) class Item(BaseModel): - prompt: str + prompt: str = "你好" history: list[tuple[str, str]] = [[]] max_length: int = 2048 top_p: float = 0.7 From be2c943b61807b4e8c84cc12501eee493e001203 Mon Sep 17 00:00:00 2001 From: Whitroom <1062015905@qq.com> Date: Wed, 12 Apr 2023 15:34:29 +0800 Subject: [PATCH 5/8] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E9=BB=98=E8=AE=A4?= =?UTF-8?q?=E5=80=BC=EF=BC=8C=E4=BB=A5=E4=BE=BF=E6=8E=A5=E5=8F=A3=E6=96=87?= =?UTF-8?q?=E6=A1=A3=E7=9B=B4=E6=8E=A5=E8=B0=83=E7=94=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api.py b/api.py index d706109..e827c53 100644 --- a/api.py +++ b/api.py @@ -37,7 +37,7 @@ app = FastAPI(lifespan=lifespan) class Item(BaseModel): prompt: str = "你好" - history: list[tuple[str, str]] = [[]] + history: list[tuple[str, str]] = [] max_length: int = 2048 top_p: float = 0.7 temperature: float = 0.95 From 14f73a2efe1278031f71cef385ae4939aff47ba9 Mon Sep 17 00:00:00 2001 From: Whitroom <1062015905@qq.com> Date: Wed, 12 Apr 2023 15:35:57 +0800 Subject: [PATCH 6/8] =?UTF-8?q?=E5=9B=9E=E8=B0=83=E7=AB=AF=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api.py b/api.py index e827c53..fd74458 100644 --- a/api.py +++ b/api.py @@ -63,4 +63,4 @@ async def create_item(item: Item): return Answer(response=response, history=history, status=200, time=time) if __name__ == '__main__': - uvicorn.run(app, host='0.0.0.0', port=8010, workers=1) + uvicorn.run(app, host='0.0.0.0', port=8000, workers=1) From d2aef0b8ad021d5be3fa5191557127913df42275 Mon Sep 17 00:00:00 2001 From: Whitroom <1062015905@qq.com> Date: Wed, 12 Apr 2023 16:04:47 +0800 Subject: [PATCH 7/8] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E5=8D=95=E5=85=83?= =?UTF-8?q?=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api_test.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 api_test.py diff --git a/api_test.py b/api_test.py new file mode 100644 index 0000000..f47bf1f --- /dev/null +++ b/api_test.py @@ -0,0 +1,17 @@ +import unittest +from httpx import AsyncClient + +class TestGenerateChat(unittest.IsolatedAsyncioTestCase): + async def test_generate_chat(self): + async with AsyncClient() as ac: + response = await ac.post( + "http://localhost:8000/", + json={ + "prompt": "你好", + "history": [], + "max_length": 2048, + "top_p": 0.7, + "temperature": 0.95 + }) + self.assertEqual(response.status_code, 200) + print(response.json()) \ No newline at end of file From 08acab314595017ee75c487c2f6be9fabaa45a41 Mon Sep 17 00:00:00 2001 From: Whitroom <1062015905@qq.com> Date: Wed, 12 Apr 2023 16:08:21 +0800 Subject: [PATCH 8/8] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E5=8D=95=E5=85=83?= =?UTF-8?q?=E6=B5=8B=E8=AF=95=E7=9A=84=E6=B3=A8=E9=87=8A?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api_test.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/api_test.py b/api_test.py index f47bf1f..1db6cbd 100644 --- a/api_test.py +++ b/api_test.py @@ -2,6 +2,17 @@ import unittest from httpx import AsyncClient class TestGenerateChat(unittest.IsolatedAsyncioTestCase): + """ + 测试生成聊天内容 + 1. 先启动服务 + ```bash + python api.py + ``` + 2. 运行测试 + ```bash + python -m unittest api_test.py + ``` + """ async def test_generate_chat(self): async with AsyncClient() as ac: response = await ac.post(