From 3ad3626176ff43f60a1f02bc0c43f6a6317b7239 Mon Sep 17 00:00:00 2001 From: Tran Thanh Luan <92072154+toilaluan@users.noreply.github.com> Date: Tue, 23 Jul 2024 16:17:14 +0700 Subject: [PATCH 1/8] add text completion --- app.py | 35 +++++++++++++++++++++++++++++++++++ requirements.txt | 3 ++- 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/app.py b/app.py index b1cb3ec..6df3fe8 100644 --- a/app.py +++ b/app.py @@ -21,6 +21,7 @@ from slowapi.util import get_remote_address from prometheus_fastapi_instrumentator import Instrumentator from PIL import Image +from transformers import AutoTokenizer def get_api_key(request: Request): return request.headers.get("API_KEY", get_remote_address(request)) @@ -83,6 +84,14 @@ class ValidatorInfo(BaseModel): all_uid_info: dict = {} sha: str = "" +class ChatCompletion(BaseModel): + model: str + messages: list[dict] + temperature: float = 1 + top_p: float = 1 + max_tokens: int = 128 + + class ImageGenerationService: def __init__(self): self.subtensor = bt.subtensor("finney") @@ -126,6 +135,12 @@ def __init__(self): Thread(target=self.sync_metagraph_periodically, daemon=True).start() Thread(target=self.recheck_validators, daemon=True).start() Thread(target=self.update_model_config, daemon=True).start() + self.tokenizer_config = self.model_config.find_one({"name": "tokenizer"}) + print(self.tokenizer_config, flush=True) + self.tokenizers = { + k, AutoTokenizer.from_pretrained(v) for k, v in self.tokenizer_config["data"] + } + print(self.tokenizers, flush=True) def update_model_config(self): while True: @@ -568,6 +583,26 @@ async def controlnet_api(self, request: Request, data: ImageToImage): generate_data["pipeline_params"][key] = value return await self.generate(Prompt(**generate_data)) + + async def chat_completions(self, request: Request, data: ChatCompletion): + # Get API_KEY from header + api_key = request.headers.get("API_KEY") + self.check_auth(api_key) + if data.model not in self.model_list: + raise HTTPException(status_code=404, detail="Model not found") + messages_str = self.tokenizers[data.model].apply_chat_template(data.messages, tokenize=False) + print(f"Chat message str: {messages_str}", flush=True) + generate_data = { + "key": api_key, + "prompt_input": messages_str, + "model_name": data.model, + "pipeline_params": { + data.temperature, + data.top_p, + data.max_tokens + } + } + return await self.generate(TextPrompt(**generate_data)) def base64_to_pil_image(self, base64_image): image = base64.b64decode(base64_image) diff --git a/requirements.txt b/requirements.txt index 259b400..d2fc359 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,5 @@ tqdm==4.66.1 httpx==0.27.0 prometheus_fastapi_instrumentator==6.0.0 pymongo==4.7.3 -slowapi==0.1.9 \ No newline at end of file +slowapi==0.1.9 +transformers \ No newline at end of file From 9563cae24c4ce3e8e5869e6b70b19a071f55bd57 Mon Sep 17 00:00:00 2001 From: Tran Thanh Luan <92072154+toilaluan@users.noreply.github.com> Date: Tue, 23 Jul 2024 16:19:21 +0700 Subject: [PATCH 2/8] fix --- app.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app.py b/app.py index 6df3fe8..6b70e9c 100644 --- a/app.py +++ b/app.py @@ -138,7 +138,7 @@ def __init__(self): self.tokenizer_config = self.model_config.find_one({"name": "tokenizer"}) print(self.tokenizer_config, flush=True) self.tokenizers = { - k, AutoTokenizer.from_pretrained(v) for k, v in self.tokenizer_config["data"] + k, AutoTokenizer.from_pretrained(v) for k, v in self.tokenizer_config["data"].items() } print(self.tokenizers, flush=True) From d4a2531910daae553d2c53bf334f4dd3fa6d076c Mon Sep 17 00:00:00 2001 From: Tran Thanh Luan <92072154+toilaluan@users.noreply.github.com> Date: Tue, 23 Jul 2024 16:20:50 +0700 Subject: [PATCH 3/8] uopdate --- app.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/app.py b/app.py index 6b70e9c..4312670 100644 --- a/app.py +++ b/app.py @@ -137,8 +137,9 @@ def __init__(self): Thread(target=self.update_model_config, daemon=True).start() self.tokenizer_config = self.model_config.find_one({"name": "tokenizer"}) print(self.tokenizer_config, flush=True) + tokenizers = [AutoTokenizer.from_pretrained(v) for v in self.tokenizer_config['data'].values()] self.tokenizers = { - k, AutoTokenizer.from_pretrained(v) for k, v in self.tokenizer_config["data"].items() + name, tokenizer for name in zip(list(self.tokenizer_config['data'].keys(), tokenizers)) } print(self.tokenizers, flush=True) From c182df3b825652f6946af45e537d4fca148866f2 Mon Sep 17 00:00:00 2001 From: Tran Thanh Luan <92072154+toilaluan@users.noreply.github.com> Date: Tue, 23 Jul 2024 16:21:07 +0700 Subject: [PATCH 4/8] fix --- app.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app.py b/app.py index 4312670..af990ce 100644 --- a/app.py +++ b/app.py @@ -139,7 +139,7 @@ def __init__(self): print(self.tokenizer_config, flush=True) tokenizers = [AutoTokenizer.from_pretrained(v) for v in self.tokenizer_config['data'].values()] self.tokenizers = { - name, tokenizer for name in zip(list(self.tokenizer_config['data'].keys(), tokenizers)) + name: tokenizer for name in zip(list(self.tokenizer_config['data'].keys(), tokenizers)) } print(self.tokenizers, flush=True) From 49af7a686db1cd76cd7dc1a42486021148d1f9b9 Mon Sep 17 00:00:00 2001 From: Tran Thanh Luan <92072154+toilaluan@users.noreply.github.com> Date: Tue, 23 Jul 2024 16:22:49 +0700 Subject: [PATCH 5/8] update --- app.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/app.py b/app.py index af990ce..1368fe6 100644 --- a/app.py +++ b/app.py @@ -137,9 +137,8 @@ def __init__(self): Thread(target=self.update_model_config, daemon=True).start() self.tokenizer_config = self.model_config.find_one({"name": "tokenizer"}) print(self.tokenizer_config, flush=True) - tokenizers = [AutoTokenizer.from_pretrained(v) for v in self.tokenizer_config['data'].values()] self.tokenizers = { - name: tokenizer for name in zip(list(self.tokenizer_config['data'].keys(), tokenizers)) + k: AutoTokenizer.from_pretrained(v) for k, v in self.tokenizer_config["data"].items() } print(self.tokenizers, flush=True) From cd6c6e367a15901f700a3e22ad836337c5d1a0b2 Mon Sep 17 00:00:00 2001 From: Tran Thanh Luan <92072154+toilaluan@users.noreply.github.com> Date: Tue, 23 Jul 2024 16:26:54 +0700 Subject: [PATCH 6/8] update --- app.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/app.py b/app.py index 1368fe6..f1b0611 100644 --- a/app.py +++ b/app.py @@ -667,3 +667,8 @@ async def instantid_api(request: Request, data: ImageToImage): @limiter.limit(API_RATE_LIMIT) # Update the rate limit async def controlnet_api(request: Request, data: ImageToImage): return await app.controlnet_api(request, data) + +@app.app.post("/api/v1/chat/completions") +@limiter.limit(API_RATE_LIMIT) +async def chat_completions_api(request: Request, data: ChatCompletion): + return await app.chat_completions(request, data) From 546d828d13d3abe49925fda4fa354c1de88f96ef Mon Sep 17 00:00:00 2001 From: Tran Thanh Luan <92072154+toilaluan@users.noreply.github.com> Date: Tue, 23 Jul 2024 16:35:28 +0700 Subject: [PATCH 7/8] update --- app.py | 6 +++--- requirements.txt | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/app.py b/app.py index f1b0611..cb82eab 100644 --- a/app.py +++ b/app.py @@ -597,9 +597,9 @@ async def chat_completions(self, request: Request, data: ChatCompletion): "prompt_input": messages_str, "model_name": data.model, "pipeline_params": { - data.temperature, - data.top_p, - data.max_tokens + "temperature": data.temperature, + "top_p": data.top_p, + "max_tokens": data.max_tokens } } return await self.generate(TextPrompt(**generate_data)) diff --git a/requirements.txt b/requirements.txt index d2fc359..94eaeff 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,5 @@ httpx==0.27.0 prometheus_fastapi_instrumentator==6.0.0 pymongo==4.7.3 slowapi==0.1.9 -transformers \ No newline at end of file +transformers +jinja2==3.1.0 \ No newline at end of file From d3850e4ba550ad96b1b36f764ecd4753a7d94806 Mon Sep 17 00:00:00 2001 From: Tran Thanh Luan <92072154+toilaluan@users.noreply.github.com> Date: Tue, 23 Jul 2024 21:33:50 +0700 Subject: [PATCH 8/8] Update app.py --- app.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/app.py b/app.py index cb82eab..3258da7 100644 --- a/app.py +++ b/app.py @@ -602,7 +602,8 @@ async def chat_completions(self, request: Request, data: ChatCompletion): "max_tokens": data.max_tokens } } - return await self.generate(TextPrompt(**generate_data)) + response = await self.generate(TextPrompt(**generate_data)) + return response['prompt_output'] def base64_to_pil_image(self, base64_image): image = base64.b64decode(base64_image)