ruslanmv commited on
Commit
6ec7105
0 Parent(s):

First commit

Browse files
.github/workflows/sync_to_hf_space.yml ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ name: Sync to Hugging Face hub
2
+ on:
3
+ push:
4
+ branches: [main]
5
+ workflow_dispatch:
6
+
7
+ jobs:
8
+ sync-to-hub:
9
+ runs-on: ubuntu-latest
10
+ steps:
11
+ - uses: actions/checkout@v3
12
+ with:
13
+ fetch-depth: 0
14
+ lfs: true
15
+ - name: Push to hub
16
+ env:
17
+ HF_TOKEN: ${{ secrets.HF_TOKEN }}
18
+ run: git push -f https://ruslanmv:$HF_TOKEN@huggingface.co/spaces/ruslanmv/hf-llm-api-collection main
.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ secrets.json
2
+ __pycache__
Dockerfile ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ FROM python:3.11-slim
2
+ WORKDIR $HOME/app
3
+ COPY . .
4
+ RUN pip install -r requirements.txt
5
+ VOLUME /data
6
+ EXPOSE 23333
7
+ CMD ["python", "-m", "apis.chat_api"]
README.md ADDED
@@ -0,0 +1,164 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: HF LLM API COLLECTION
3
+ emoji: ☯️
4
+ colorFrom: gray
5
+ colorTo: gray
6
+ sdk: docker
7
+ app_port: 23333
8
+ ---
9
+
10
+ ## HF-LLM-API
11
+ Huggingface LLM Inference API in OpenAI message format.
12
+
13
+ Project link: https://github.com/ruslanmv/hf-llm-api-collection
14
+
15
+ ## Features
16
+
17
+ - Available Models (2024/01/22): [#5](https://github.com/Hansimov/hf-llm-api/issues/5)
18
+ - `mistral-7b`, `mixtral-8x7b`, `nous-mixtral-8x7b`
19
+ - Adaptive prompt templates for different models
20
+ - Support OpenAI API format
21
+ - Enable api endpoint via official `openai-python` package
22
+ - Support both stream and no-stream response
23
+ - Support API Key via both HTTP auth header and env varible [#4](https://github.com/Hansimov/hf-llm-api/issues/4)
24
+ - Docker deployment
25
+
26
+ ## Run API service
27
+
28
+ ### Run in Command Line
29
+
30
+ **Install dependencies:**
31
+
32
+ ```bash
33
+ # pipreqs . --force --mode no-pin
34
+ pip install -r requirements.txt
35
+ ```
36
+
37
+ **Run API:**
38
+
39
+ ```bash
40
+ python -m apis.chat_api
41
+ ```
42
+
43
+ ## Run via Docker
44
+
45
+ **Docker build:**
46
+
47
+ ```bash
48
+ sudo docker build -t hf-llm-api:1.0 . --build-arg http_proxy=$http_proxy --build-arg https_proxy=$https_proxy
49
+ ```
50
+
51
+ **Docker run:**
52
+
53
+ ```bash
54
+ # no proxy
55
+ sudo docker run -p 23333:23333 hf-llm-api:1.0
56
+
57
+ # with proxy
58
+ sudo docker run -p 23333:23333 --env http_proxy="http://<server>:<port>" hf-llm-api:1.0
59
+ ```
60
+
61
+ ## API Usage
62
+
63
+ ### Using `openai-python`
64
+
65
+ See: [`examples/chat_with_openai.py`](https://github.com/ruslanmv/hf-llm-api-collection/blob/main/examples/chat_with_openai.py)
66
+
67
+ ```py
68
+ from openai import OpenAI
69
+
70
+ # If runnning this service with proxy, you might need to unset `http(s)_proxy`.
71
+ base_url = "http://127.0.0.1:23333"
72
+ # Your own HF_TOKEN
73
+ api_key = "hf_xxxxxxxxxxxxxxxx"
74
+ # use below as non-auth user
75
+ # api_key = "sk-xxx"
76
+
77
+ client = OpenAI(base_url=base_url, api_key=api_key)
78
+ response = client.chat.completions.create(
79
+ model="mixtral-8x7b",
80
+ messages=[
81
+ {
82
+ "role": "user",
83
+ "content": "what is your model",
84
+ }
85
+ ],
86
+ stream=True,
87
+ )
88
+
89
+ for chunk in response:
90
+ if chunk.choices[0].delta.content is not None:
91
+ print(chunk.choices[0].delta.content, end="", flush=True)
92
+ elif chunk.choices[0].finish_reason == "stop":
93
+ print()
94
+ else:
95
+ pass
96
+ ```
97
+
98
+ ### Using post requests
99
+
100
+ See: [`examples/chat_with_post.py`](https://github.com/ruslanmv/hf-llm-api-collection/blob/main/examples/chat_with_post.py)
101
+
102
+
103
+ ```py
104
+ import ast
105
+ import httpx
106
+ import json
107
+ import re
108
+
109
+ # If runnning this service with proxy, you might need to unset `http(s)_proxy`.
110
+ chat_api = "http://127.0.0.1:23333"
111
+ # Your own HF_TOKEN
112
+ api_key = "hf_xxxxxxxxxxxxxxxx"
113
+ # use below as non-auth user
114
+ # api_key = "sk-xxx"
115
+
116
+ requests_headers = {}
117
+ requests_payload = {
118
+ "model": "mixtral-8x7b",
119
+ "messages": [
120
+ {
121
+ "role": "user",
122
+ "content": "what is your model",
123
+ }
124
+ ],
125
+ "stream": True,
126
+ }
127
+
128
+ with httpx.stream(
129
+ "POST",
130
+ chat_api + "/chat/completions",
131
+ headers=requests_headers,
132
+ json=requests_payload,
133
+ timeout=httpx.Timeout(connect=20, read=60, write=20, pool=None),
134
+ ) as response:
135
+ # https://docs.aiohttp.org/en/stable/streams.html
136
+ # https://github.com/openai/openai-cookbook/blob/main/examples/How_to_stream_completions.ipynb
137
+ response_content = ""
138
+ for line in response.iter_lines():
139
+ remove_patterns = [r"^\s*data:\s*", r"^\s*\[DONE\]\s*"]
140
+ for pattern in remove_patterns:
141
+ line = re.sub(pattern, "", line).strip()
142
+
143
+ if line:
144
+ try:
145
+ line_data = json.loads(line)
146
+ except Exception as e:
147
+ try:
148
+ line_data = ast.literal_eval(line)
149
+ except:
150
+ print(f"Error: {line}")
151
+ raise e
152
+ # print(f"line: {line_data}")
153
+ delta_data = line_data["choices"][0]["delta"]
154
+ finish_reason = line_data["choices"][0]["finish_reason"]
155
+ if "role" in delta_data:
156
+ role = delta_data["role"]
157
+ if "content" in delta_data:
158
+ delta_content = delta_data["content"]
159
+ response_content += delta_content
160
+ print(delta_content, end="", flush=True)
161
+ if finish_reason == "stop":
162
+ print()
163
+
164
+ ```
__init__.py ADDED
File without changes
apis/__init__.py ADDED
File without changes
apis/chat_api.py ADDED
@@ -0,0 +1,214 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import markdown2
3
+ import os
4
+ import sys
5
+ import uvicorn
6
+
7
+ from pathlib import Path
8
+ from fastapi import FastAPI, Depends
9
+ from fastapi.responses import HTMLResponse
10
+ from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
11
+ from pydantic import BaseModel, Field
12
+ from typing import Union
13
+ from sse_starlette.sse import EventSourceResponse, ServerSentEvent
14
+ from utils.logger import logger
15
+ from networks.message_streamer import MessageStreamer
16
+ from messagers.message_composer import MessageComposer
17
+ from mocks.stream_chat_mocker import stream_chat_mock
18
+
19
+
20
+ class ChatAPIApp:
21
+ def __init__(self):
22
+ self.app = FastAPI(
23
+ docs_url="/",
24
+ title="HuggingFace LLM API",
25
+ swagger_ui_parameters={"defaultModelsExpandDepth": -1},
26
+ version="1.0",
27
+ )
28
+ self.setup_routes()
29
+
30
+ def get_available_models(self):
31
+ # https://platform.openai.com/docs/api-reference/models/list
32
+ # ANCHOR[id=available-models]: Available models
33
+ self.available_models = {
34
+ "object": "list",
35
+ "data": [
36
+ {
37
+ "id": "mixtral-8x7b",
38
+ "description": "[mistralai/Mixtral-8x7B-Instruct-v0.1]: https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1",
39
+ "object": "model",
40
+ "created": 1700000000,
41
+ "owned_by": "mistralai",
42
+ },
43
+ {
44
+ "id": "mistral-7b",
45
+ "description": "[mistralai/Mistral-7B-Instruct-v0.2]: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2",
46
+ "object": "model",
47
+ "created": 1700000000,
48
+ "owned_by": "mistralai",
49
+ },
50
+ {
51
+ "id": "nous-mixtral-8x7b",
52
+ "description": "[NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO]: https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
53
+ "object": "model",
54
+ "created": 1700000000,
55
+ "owned_by": "NousResearch",
56
+ },
57
+ ],
58
+ }
59
+ return self.available_models
60
+
61
+ def extract_api_key(
62
+ credentials: HTTPAuthorizationCredentials = Depends(
63
+ HTTPBearer(auto_error=False)
64
+ ),
65
+ ):
66
+ api_key = None
67
+ if credentials:
68
+ api_key = credentials.credentials
69
+ else:
70
+ api_key = os.getenv("HF_TOKEN")
71
+
72
+ if api_key:
73
+ if api_key.startswith("hf_"):
74
+ return api_key
75
+ else:
76
+ logger.warn(f"Invalid HF Token!")
77
+ else:
78
+ logger.warn("Not provide HF Token!")
79
+ return None
80
+
81
+ class ChatCompletionsPostItem(BaseModel):
82
+ model: str = Field(
83
+ default="mixtral-8x7b",
84
+ description="(str) `mixtral-8x7b`",
85
+ )
86
+ messages: list = Field(
87
+ default=[{"role": "user", "content": "Hello, who are you?"}],
88
+ description="(list) Messages",
89
+ )
90
+ temperature: Union[float, None] = Field(
91
+ default=0.5,
92
+ description="(float) Temperature",
93
+ )
94
+ top_p: Union[float, None] = Field(
95
+ default=0.95,
96
+ description="(float) top p",
97
+ )
98
+ max_tokens: Union[int, None] = Field(
99
+ default=-1,
100
+ description="(int) Max tokens",
101
+ )
102
+ use_cache: bool = Field(
103
+ default=False,
104
+ description="(bool) Use cache",
105
+ )
106
+ stream: bool = Field(
107
+ default=True,
108
+ description="(bool) Stream",
109
+ )
110
+
111
+ def chat_completions(
112
+ self, item: ChatCompletionsPostItem, api_key: str = Depends(extract_api_key)
113
+ ):
114
+ streamer = MessageStreamer(model=item.model)
115
+ composer = MessageComposer(model=item.model)
116
+ composer.merge(messages=item.messages)
117
+ # streamer.chat = stream_chat_mock
118
+
119
+ stream_response = streamer.chat_response(
120
+ prompt=composer.merged_str,
121
+ temperature=item.temperature,
122
+ top_p=item.top_p,
123
+ max_new_tokens=item.max_tokens,
124
+ api_key=api_key,
125
+ use_cache=item.use_cache,
126
+ )
127
+ if item.stream:
128
+ event_source_response = EventSourceResponse(
129
+ streamer.chat_return_generator(stream_response),
130
+ media_type="text/event-stream",
131
+ ping=2000,
132
+ ping_message_factory=lambda: ServerSentEvent(**{"comment": ""}),
133
+ )
134
+ return event_source_response
135
+ else:
136
+ data_response = streamer.chat_return_dict(stream_response)
137
+ return data_response
138
+
139
+ def get_readme(self):
140
+ readme_path = Path(__file__).parents[1] / "README.md"
141
+ with open(readme_path, "r", encoding="utf-8") as rf:
142
+ readme_str = rf.read()
143
+ readme_html = markdown2.markdown(
144
+ readme_str, extras=["table", "fenced-code-blocks", "highlightjs-lang"]
145
+ )
146
+ return readme_html
147
+
148
+ def setup_routes(self):
149
+ for prefix in ["", "/v1", "/api", "/api/v1"]:
150
+ if prefix in ["/api/v1"]:
151
+ include_in_schema = True
152
+ else:
153
+ include_in_schema = False
154
+
155
+ self.app.get(
156
+ prefix + "/models",
157
+ summary="Get available models",
158
+ include_in_schema=include_in_schema,
159
+ )(self.get_available_models)
160
+
161
+ self.app.post(
162
+ prefix + "/chat/completions",
163
+ summary="Chat completions in conversation session",
164
+ include_in_schema=include_in_schema,
165
+ )(self.chat_completions)
166
+ self.app.get(
167
+ "/readme",
168
+ summary="README of HF LLM API",
169
+ response_class=HTMLResponse,
170
+ include_in_schema=False,
171
+ )(self.get_readme)
172
+
173
+
174
+ class ArgParser(argparse.ArgumentParser):
175
+ def __init__(self, *args, **kwargs):
176
+ super(ArgParser, self).__init__(*args, **kwargs)
177
+
178
+ self.add_argument(
179
+ "-s",
180
+ "--server",
181
+ type=str,
182
+ default="0.0.0.0",
183
+ help="Server IP for HF LLM Chat API",
184
+ )
185
+ self.add_argument(
186
+ "-p",
187
+ "--port",
188
+ type=int,
189
+ default=23333,
190
+ help="Server Port for HF LLM Chat API",
191
+ )
192
+
193
+ self.add_argument(
194
+ "-d",
195
+ "--dev",
196
+ default=False,
197
+ action="store_true",
198
+ help="Run in dev mode",
199
+ )
200
+
201
+ self.args = self.parse_args(sys.argv[1:])
202
+
203
+
204
+ app = ChatAPIApp().app
205
+
206
+ if __name__ == "__main__":
207
+ args = ArgParser().args
208
+ if args.dev:
209
+ uvicorn.run("__main__:app", host=args.server, port=args.port, reload=True)
210
+ else:
211
+ uvicorn.run("__main__:app", host=args.server, port=args.port, reload=False)
212
+
213
+ # python -m apis.chat_api # [Docker] on product mode
214
+ # python -m apis.chat_api -d # [Dev] on develop mode
examples/__init__.py ADDED
File without changes
examples/chat_with_openai.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from openai import OpenAI
2
+
3
+ # If runnning this service with proxy, you might need to unset `http(s)_proxy`.
4
+ base_url = "http://127.0.0.1:23333"
5
+ api_key = "sk-xxxxx"
6
+
7
+ client = OpenAI(base_url=base_url, api_key=api_key)
8
+ response = client.chat.completions.create(
9
+ model="mixtral-8x7b",
10
+ messages=[
11
+ {
12
+ "role": "user",
13
+ "content": "what is your model",
14
+ }
15
+ ],
16
+ stream=True,
17
+ )
18
+
19
+ for chunk in response:
20
+ if chunk.choices[0].delta.content is not None:
21
+ print(chunk.choices[0].delta.content, end="", flush=True)
22
+ elif chunk.choices[0].finish_reason == "stop":
23
+ print()
24
+ else:
25
+ pass
examples/chat_with_post.py ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import ast
2
+ import httpx
3
+ import json
4
+ import re
5
+
6
+ # If runnning this service with proxy, you might need to unset `http(s)_proxy`.
7
+ chat_api = "http://127.0.0.1:23333"
8
+ api_key = "sk-xxxxx"
9
+ requests_headers = {}
10
+ requests_payload = {
11
+ "model": "mixtral-8x7b",
12
+ "messages": [
13
+ {
14
+ "role": "user",
15
+ "content": "what is your model",
16
+ }
17
+ ],
18
+ "stream": True,
19
+ }
20
+
21
+ with httpx.stream(
22
+ "POST",
23
+ chat_api + "/chat/completions",
24
+ headers=requests_headers,
25
+ json=requests_payload,
26
+ timeout=httpx.Timeout(connect=20, read=60, write=20, pool=None),
27
+ ) as response:
28
+ # https://docs.aiohttp.org/en/stable/streams.html
29
+ # https://github.com/openai/openai-cookbook/blob/main/examples/How_to_stream_completions.ipynb
30
+ response_content = ""
31
+ for line in response.iter_lines():
32
+ remove_patterns = [r"^\s*data:\s*", r"^\s*\[DONE\]\s*"]
33
+ for pattern in remove_patterns:
34
+ line = re.sub(pattern, "", line).strip()
35
+
36
+ if line:
37
+ try:
38
+ line_data = json.loads(line)
39
+ except Exception as e:
40
+ try:
41
+ line_data = ast.literal_eval(line)
42
+ except:
43
+ print(f"Error: {line}")
44
+ raise e
45
+ # print(f"line: {line_data}")
46
+ delta_data = line_data["choices"][0]["delta"]
47
+ finish_reason = line_data["choices"][0]["finish_reason"]
48
+ if "role" in delta_data:
49
+ role = delta_data["role"]
50
+ if "content" in delta_data:
51
+ delta_content = delta_data["content"]
52
+ response_content += delta_content
53
+ print(delta_content, end="", flush=True)
54
+ if finish_reason == "stop":
55
+ print()
messagers/__init__.py ADDED
File without changes
messagers/message_composer.py ADDED
@@ -0,0 +1,247 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ from pprint import pprint
3
+ from utils.logger import logger
4
+
5
+
6
+ class MessageComposer:
7
+ # LINK - apis/chat_api.py#available-models
8
+ AVALAIBLE_MODELS = [
9
+ "mixtral-8x7b",
10
+ "mistral-7b",
11
+ "openchat-3.5",
12
+ "nous-mixtral-8x7b",
13
+ ]
14
+
15
+ def __init__(self, model: str = None):
16
+ if model in self.AVALAIBLE_MODELS:
17
+ self.model = model
18
+ else:
19
+ self.model = "mixtral-8x7b"
20
+ self.system_roles = ["system"]
21
+ self.inst_roles = ["user", "system", "inst"]
22
+ self.answer_roles = ["assistant", "bot", "answer"]
23
+ self.default_role = "user"
24
+
25
+ def concat_messages_by_role(self, messages):
26
+ def is_same_role(role1, role2):
27
+ if (
28
+ (role1 == role2)
29
+ or (role1 in self.inst_roles and role2 in self.inst_roles)
30
+ or (role1 in self.answer_roles and role2 in self.answer_roles)
31
+ ):
32
+ return True
33
+ else:
34
+ return False
35
+
36
+ concat_messages = []
37
+ for message in messages:
38
+ role = message["role"]
39
+ content = message["content"]
40
+ if concat_messages and is_same_role(role, concat_messages[-1]["role"]):
41
+ concat_messages[-1]["content"] += "\n" + content
42
+ else:
43
+ if role in self.inst_roles:
44
+ message["role"] = "inst"
45
+ elif role in self.answer_roles:
46
+ message["role"] = "answer"
47
+ else:
48
+ message["role"] = "inst"
49
+ concat_messages.append(message)
50
+ return concat_messages
51
+
52
+ def merge(self, messages) -> str:
53
+ # Mistral and Mixtral:
54
+ # <s> [INST] Instruction [/INST] Model answer </s> [INST] Follow-up instruction [/INST]
55
+
56
+ # OpenChat:
57
+ # GPT4 Correct User: Hello<|end_of_turn|>GPT4 Correct Assistant: Hi<|end_of_turn|>GPT4 Correct User: How are you today?<|end_of_turn|>GPT4 Correct Assistant:
58
+
59
+ # Nous Mixtral:
60
+ # <|im_start|>system
61
+ # You are "Hermes 2".<|im_end|>
62
+ # <|im_start|>user
63
+ # Hello, who are you?<|im_end|>
64
+ # <|im_start|>assistant
65
+
66
+ self.messages = messages
67
+ self.merged_str = ""
68
+
69
+ # https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1#instruction-format
70
+ if self.model in ["mixtral-8x7b", "mistral-7b"]:
71
+ self.messages = self.concat_messages_by_role(messages)
72
+ self.cached_str = ""
73
+ for message in self.messages:
74
+ role = message["role"]
75
+ content = message["content"]
76
+ if role in self.inst_roles:
77
+ self.cached_str = f"[INST] {content} [/INST]"
78
+ elif role in self.answer_roles:
79
+ self.merged_str += f"<s> {self.cached_str} {content} </s>\n"
80
+ self.cached_str = ""
81
+ else:
82
+ self.cached_str = f"[INST] {content} [/INST]"
83
+ if self.cached_str:
84
+ self.merged_str += f"{self.cached_str}"
85
+ # https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO#prompt-format
86
+ elif self.model in ["nous-mixtral-8x7b"]:
87
+ self.merged_str_list = []
88
+ for message in self.messages:
89
+ role = message["role"]
90
+ content = message["content"]
91
+ if role not in ["system", "user", "assistant"]:
92
+ role = self.default_role
93
+ message_line = f"<|im_start|>{role}\n{content}<|im_end|>"
94
+ self.merged_str_list.append(message_line)
95
+ self.merged_str_list.append("<|im_start|>assistant")
96
+ self.merged_str = "\n".join(self.merged_str_list)
97
+ # https://huggingface.co/openchat/openchat-3.5-0106
98
+ elif self.model in ["openchat-3.5"]:
99
+ self.messages = self.concat_messages_by_role(messages)
100
+ self.merged_str_list = []
101
+ self.end_of_turn = "<|end_of_turn|>"
102
+ for message in self.messages:
103
+ role = message["role"]
104
+ content = message["content"]
105
+ if role in self.inst_roles:
106
+ self.merged_str_list.append(
107
+ f"GPT4 Correct User:\n{content}{self.end_of_turn}"
108
+ )
109
+ elif role in self.answer_roles:
110
+ self.merged_str_list.append(
111
+ f"GPT4 Correct Assistant:\n{content}{self.end_of_turn}"
112
+ )
113
+ else:
114
+ self.merged_str_list.append(
115
+ f"GPT4 Correct User: {content}{self.end_of_turn}"
116
+ )
117
+ self.merged_str_list.append(f"GPT4 Correct Assistant:\n")
118
+ self.merged_str = "\n".join(self.merged_str_list)
119
+ else:
120
+ self.merged_str = "\n".join(
121
+ [
122
+ f'`{message["role"]}`:\n{message["content"]}\n'
123
+ for message in self.messages
124
+ ]
125
+ )
126
+
127
+ return self.merged_str
128
+
129
+ def convert_pair_matches_to_messages(self, pair_matches_list):
130
+ messages = []
131
+ if len(pair_matches_list) <= 0:
132
+ messages = [
133
+ {
134
+ "role": "user",
135
+ "content": self.merged_str,
136
+ }
137
+ ]
138
+ else:
139
+ for match in pair_matches_list:
140
+ inst = match.group("inst")
141
+ answer = match.group("answer")
142
+ messages.extend(
143
+ [
144
+ {"role": "user", "content": inst.strip()},
145
+ {"role": "assistant", "content": answer.strip()},
146
+ ]
147
+ )
148
+ return messages
149
+
150
+ def append_last_instruction_to_messages(self, inst_matches_list, pair_matches_list):
151
+ if len(inst_matches_list) > len(pair_matches_list):
152
+ self.messages.extend(
153
+ [
154
+ {
155
+ "role": "user",
156
+ "content": inst_matches_list[-1].group("inst").strip(),
157
+ }
158
+ ]
159
+ )
160
+
161
+ def split(self, merged_str) -> list:
162
+ self.merged_str = merged_str
163
+ self.messages = []
164
+
165
+ if self.model in ["mixtral-8x7b", "mistral-7b"]:
166
+ pair_pattern = (
167
+ r"<s>\s*\[INST\](?P<inst>[\s\S]*?)\[/INST\](?P<answer>[\s\S]*?)</s>"
168
+ )
169
+ pair_matches = re.finditer(pair_pattern, self.merged_str, re.MULTILINE)
170
+ pair_matches_list = list(pair_matches)
171
+
172
+ self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
173
+
174
+ inst_pattern = r"\[INST\](?P<inst>[\s\S]*?)\[/INST\]"
175
+ inst_matches = re.finditer(inst_pattern, self.merged_str, re.MULTILINE)
176
+ inst_matches_list = list(inst_matches)
177
+
178
+ self.append_last_instruction_to_messages(
179
+ inst_matches_list, pair_matches_list
180
+ )
181
+ elif self.model in ["nous-mixtral-8x7b"]:
182
+ # https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO#prompt-format
183
+ # message_pattern = r"<\|im_start\|>(?P<role>system|user|assistant)[\s\n]*(?P<content>[\s\S]*?)<\|im_end\|>"
184
+ message_pattern = r"<\|im_start\|>(?P<role>system|user|assistant)[\s\n]*(?P<content>[\s\S]*?)<\|im_end\|>"
185
+ message_matches = re.finditer(
186
+ message_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
187
+ )
188
+ message_matches_list = list(message_matches)
189
+ logger.note(f"message_matches_list: {message_matches_list}")
190
+ for match in message_matches_list:
191
+ role = match.group("role")
192
+ content = match.group("content")
193
+ self.messages.append({"role": role, "content": content.strip()})
194
+ elif self.model in ["openchat-3.5"]:
195
+ pair_pattern = r"GPT4 Correct User:(?P<inst>[\s\S]*?)<\|end_of_turn\|>\s*GPT4 Correct Assistant:(?P<answer>[\s\S]*?)<\|end_of_turn\|>"
196
+ pair_matches = re.finditer(
197
+ pair_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
198
+ )
199
+ pair_matches_list = list(pair_matches)
200
+ self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
201
+ inst_pattern = r"GPT4 Correct User:(?P<inst>[\s\S]*?)<\|end_of_turn\|>"
202
+ inst_matches = re.finditer(
203
+ inst_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
204
+ )
205
+ inst_matches_list = list(inst_matches)
206
+ self.append_last_instruction_to_messages(
207
+ inst_matches_list, pair_matches_list
208
+ )
209
+ else:
210
+ self.messages = [
211
+ {
212
+ "role": "user",
213
+ "content": self.merged_str,
214
+ }
215
+ ]
216
+
217
+ return self.messages
218
+
219
+
220
+ if __name__ == "__main__":
221
+ model = "mixtral-8x7b"
222
+ # model = "nous-mixtral-8x7b"
223
+ composer = MessageComposer(model)
224
+ messages = [
225
+ {
226
+ "role": "system",
227
+ "content": "You are a LLM developed by OpenAI.\nYour name is GPT-4.",
228
+ },
229
+ {"role": "user", "content": "Hello, who are you?"},
230
+ {"role": "assistant", "content": "I am a bot."},
231
+ {"role": "user", "content": "What is your name?"},
232
+ # {"role": "assistant", "content": "My name is Bing."},
233
+ # {"role": "user", "content": "Tell me a joke."},
234
+ # {"role": "assistant", "content": "What is a robot's favorite type of music?"},
235
+ # {
236
+ # "role": "user",
237
+ # "content": "How many questions have I asked? Please list them.",
238
+ # },
239
+ ]
240
+ logger.note(f"model: {composer.model}")
241
+ merged_str = composer.merge(messages)
242
+ logger.note("merged_str:")
243
+ logger.mesg(merged_str)
244
+ logger.note("splitted messages:")
245
+ pprint(composer.split(merged_str))
246
+ # logger.note("merged merged_str:")
247
+ # logger.mesg(composer.merge(composer.split(merged_str)))
messagers/message_outputer.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+
3
+
4
+ class OpenaiStreamOutputer:
5
+ """
6
+ Create chat completion - OpenAI API Documentation
7
+ * https://platform.openai.com/docs/api-reference/chat/create
8
+ """
9
+
10
+ def __init__(self):
11
+ self.default_data = {
12
+ "created": 1700000000,
13
+ "id": "chatcmpl-hugginface",
14
+ "object": "chat.completion.chunk",
15
+ # "content_type": "Completions",
16
+ "model": "hugginface",
17
+ "choices": [],
18
+ }
19
+
20
+ def data_to_string(self, data={}, content_type=""):
21
+ data_str = f"{json.dumps(data)}"
22
+ return data_str
23
+
24
+ def output(self, content=None, content_type="Completions") -> str:
25
+ data = self.default_data.copy()
26
+ if content_type == "Role":
27
+ data["choices"] = [
28
+ {
29
+ "index": 0,
30
+ "delta": {"role": "assistant"},
31
+ "finish_reason": None,
32
+ }
33
+ ]
34
+ elif content_type in [
35
+ "Completions",
36
+ "InternalSearchQuery",
37
+ "InternalSearchResult",
38
+ "SuggestedResponses",
39
+ ]:
40
+ if content_type in ["InternalSearchQuery", "InternalSearchResult"]:
41
+ content += "\n"
42
+ data["choices"] = [
43
+ {
44
+ "index": 0,
45
+ "delta": {"content": content},
46
+ "finish_reason": None,
47
+ }
48
+ ]
49
+ elif content_type == "Finished":
50
+ data["choices"] = [
51
+ {
52
+ "index": 0,
53
+ "delta": {},
54
+ "finish_reason": "stop",
55
+ }
56
+ ]
57
+ else:
58
+ data["choices"] = [
59
+ {
60
+ "index": 0,
61
+ "delta": {},
62
+ "finish_reason": None,
63
+ }
64
+ ]
65
+ return self.data_to_string(data, content_type)
mocks/__init__.py ADDED
File without changes
mocks/stream_chat_mocker.py ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import time
2
+ from utils.logger import logger
3
+
4
+
5
+ def stream_chat_mock(*args, **kwargs):
6
+ logger.note(msg=str(args) + str(kwargs))
7
+ for i in range(10):
8
+ content = f"W{i+1} "
9
+ time.sleep(0.1)
10
+ logger.mesg(content, end="")
11
+ yield content
12
+ logger.mesg("")
13
+ yield ""
networks/__init__.py ADDED
File without changes
networks/message_streamer.py ADDED
@@ -0,0 +1,212 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import re
3
+ import requests
4
+ from tiktoken import get_encoding as tiktoken_get_encoding
5
+ from messagers.message_outputer import OpenaiStreamOutputer
6
+ from utils.logger import logger
7
+ from utils.enver import enver
8
+
9
+
10
+ class MessageStreamer:
11
+ MODEL_MAP = {
12
+ "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1", # 72.62, fast [Recommended]
13
+ "mistral-7b": "mistralai/Mistral-7B-Instruct-v0.2", # 65.71, fast
14
+ "nous-mixtral-8x7b": "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
15
+ # "openchat-3.5": "openchat/openchat-3.5-1210", # 68.89, fast
16
+ # "zephyr-7b-beta": "HuggingFaceH4/zephyr-7b-beta", # ❌ Too Slow
17
+ # "llama-70b": "meta-llama/Llama-2-70b-chat-hf", # ❌ Require Pro User
18
+ # "codellama-34b": "codellama/CodeLlama-34b-Instruct-hf", # ❌ Low Score
19
+ # "falcon-180b": "tiiuae/falcon-180B-chat", # ❌ Require Pro User
20
+ "default": "mistralai/Mixtral-8x7B-Instruct-v0.1",
21
+ }
22
+ STOP_SEQUENCES_MAP = {
23
+ "mixtral-8x7b": "</s>",
24
+ "mistral-7b": "</s>",
25
+ "nous-mixtral-8x7b": "<|im_end|>",
26
+ "openchat-3.5": "<|end_of_turn|>",
27
+ }
28
+ TOKEN_LIMIT_MAP = {
29
+ "mixtral-8x7b": 32768,
30
+ "mistral-7b": 32768,
31
+ "nous-mixtral-8x7b": 32768,
32
+ "openchat-3.5": 8192,
33
+ }
34
+ TOKEN_RESERVED = 100
35
+
36
+ def __init__(self, model: str):
37
+ if model in self.MODEL_MAP.keys():
38
+ self.model = model
39
+ else:
40
+ self.model = "default"
41
+ self.model_fullname = self.MODEL_MAP[self.model]
42
+ self.message_outputer = OpenaiStreamOutputer()
43
+ self.tokenizer = tiktoken_get_encoding("cl100k_base")
44
+
45
+ def parse_line(self, line):
46
+ line = line.decode("utf-8")
47
+ line = re.sub(r"data:\s*", "", line)
48
+ data = json.loads(line)
49
+ try:
50
+ content = data["token"]["text"]
51
+ except:
52
+ logger.err(data)
53
+ return content
54
+
55
+ def count_tokens(self, text):
56
+ tokens = self.tokenizer.encode(text)
57
+ token_count = len(tokens)
58
+ logger.note(f"Prompt Token Count: {token_count}")
59
+ return token_count
60
+
61
+ def chat_response(
62
+ self,
63
+ prompt: str = None,
64
+ temperature: float = 0.5,
65
+ top_p: float = 0.95,
66
+ max_new_tokens: int = None,
67
+ api_key: str = None,
68
+ use_cache: bool = False,
69
+ ):
70
+ # https://huggingface.co/docs/api-inference/detailed_parameters?code=curl
71
+ # curl --proxy http://<server>:<port> https://api-inference.huggingface.co/models/<org>/<model_name> -X POST -d '{"inputs":"who are you?","parameters":{"max_new_token":64}}' -H 'Content-Type: application/json' -H 'Authorization: Bearer <HF_TOKEN>'
72
+ self.request_url = (
73
+ f"https://api-inference.huggingface.co/models/{self.model_fullname}"
74
+ )
75
+ self.request_headers = {
76
+ "Content-Type": "application/json",
77
+ }
78
+
79
+ if api_key:
80
+ logger.note(
81
+ f"Using API Key: {api_key[:3]}{(len(api_key)-7)*'*'}{api_key[-4:]}"
82
+ )
83
+ self.request_headers["Authorization"] = f"Bearer {api_key}"
84
+
85
+ if temperature is None or temperature < 0:
86
+ temperature = 0.0
87
+ # temperature must 0 < and < 1 for HF LLM models
88
+ temperature = max(temperature, 0.01)
89
+ temperature = min(temperature, 0.99)
90
+ top_p = max(top_p, 0.01)
91
+ top_p = min(top_p, 0.99)
92
+
93
+ token_limit = int(
94
+ self.TOKEN_LIMIT_MAP[self.model]
95
+ - self.TOKEN_RESERVED
96
+ - self.count_tokens(prompt) * 1.35
97
+ )
98
+ if token_limit <= 0:
99
+ raise ValueError("Prompt exceeded token limit!")
100
+
101
+ if max_new_tokens is None or max_new_tokens <= 0:
102
+ max_new_tokens = token_limit
103
+ else:
104
+ max_new_tokens = min(max_new_tokens, token_limit)
105
+
106
+ # References:
107
+ # huggingface_hub/inference/_client.py:
108
+ # class InferenceClient > def text_generation()
109
+ # huggingface_hub/inference/_text_generation.py:
110
+ # class TextGenerationRequest > param `stream`
111
+ # https://huggingface.co/docs/text-generation-inference/conceptual/streaming#streaming-with-curl
112
+ # https://huggingface.co/docs/api-inference/detailed_parameters#text-generation-task
113
+ self.request_body = {
114
+ "inputs": prompt,
115
+ "parameters": {
116
+ "temperature": temperature,
117
+ "top_p": top_p,
118
+ "max_new_tokens": max_new_tokens,
119
+ "return_full_text": False,
120
+ },
121
+ "options": {
122
+ "use_cache": use_cache,
123
+ },
124
+ "stream": True,
125
+ }
126
+
127
+ if self.model in self.STOP_SEQUENCES_MAP.keys():
128
+ self.stop_sequences = self.STOP_SEQUENCES_MAP[self.model]
129
+ # self.request_body["parameters"]["stop_sequences"] = [
130
+ # self.STOP_SEQUENCES[self.model]
131
+ # ]
132
+
133
+ logger.back(self.request_url)
134
+ enver.set_envs(proxies=True)
135
+ stream_response = requests.post(
136
+ self.request_url,
137
+ headers=self.request_headers,
138
+ json=self.request_body,
139
+ proxies=enver.requests_proxies,
140
+ stream=True,
141
+ )
142
+ status_code = stream_response.status_code
143
+ if status_code == 200:
144
+ logger.success(status_code)
145
+ else:
146
+ logger.err(status_code)
147
+
148
+ return stream_response
149
+
150
+ def chat_return_dict(self, stream_response):
151
+ # https://platform.openai.com/docs/guides/text-generation/chat-completions-response-format
152
+ final_output = self.message_outputer.default_data.copy()
153
+ final_output["choices"] = [
154
+ {
155
+ "index": 0,
156
+ "finish_reason": "stop",
157
+ "message": {
158
+ "role": "assistant",
159
+ "content": "",
160
+ },
161
+ }
162
+ ]
163
+ logger.back(final_output)
164
+
165
+ final_content = ""
166
+ for line in stream_response.iter_lines():
167
+ if not line:
168
+ continue
169
+ content = self.parse_line(line)
170
+
171
+ if content.strip() == self.stop_sequences:
172
+ logger.success("\n[Finished]")
173
+ break
174
+ else:
175
+ logger.back(content, end="")
176
+ final_content += content
177
+
178
+ if self.model in self.STOP_SEQUENCES_MAP.keys():
179
+ final_content = final_content.replace(self.stop_sequences, "")
180
+
181
+ final_content = final_content.strip()
182
+ final_output["choices"][0]["message"]["content"] = final_content
183
+ return final_output
184
+
185
+ def chat_return_generator(self, stream_response):
186
+ is_finished = False
187
+ line_count = 0
188
+ for line in stream_response.iter_lines():
189
+ if line:
190
+ line_count += 1
191
+ else:
192
+ continue
193
+
194
+ content = self.parse_line(line)
195
+
196
+ if content.strip() == self.stop_sequences:
197
+ content_type = "Finished"
198
+ logger.success("\n[Finished]")
199
+ is_finished = True
200
+ else:
201
+ content_type = "Completions"
202
+ if line_count == 1:
203
+ content = content.lstrip()
204
+ logger.back(content, end="")
205
+
206
+ output = self.message_outputer.output(
207
+ content=content, content_type=content_type
208
+ )
209
+ yield output
210
+
211
+ if not is_finished:
212
+ yield self.message_outputer.output(content="", content_type="Finished")
requirements.txt ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ aiohttp
2
+ fastapi
3
+ httpx
4
+ markdown2[all]
5
+ openai
6
+ pydantic
7
+ requests
8
+ sse_starlette
9
+ termcolor
10
+ tiktoken
11
+ uvicorn
12
+ websockets
utils/__init__.py ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import requests
3
+ import os
4
+
5
+ from pathlib import Path
6
+
7
+
8
+ class OSEnver:
9
+ def __init__(self):
10
+ self.envs_stack = []
11
+ self.envs = os.environ.copy()
12
+
13
+ def store_envs(self):
14
+ self.envs_stack.append(self.envs)
15
+
16
+ def restore_envs(self):
17
+ self.envs = self.envs_stack.pop()
18
+ if self.global_scope:
19
+ os.environ = self.envs
20
+
21
+ def set_envs(self, secrets=True, proxies=None, store_envs=True):
22
+ # caller_info = inspect.stack()[1]
23
+ # logger.back(f"OS Envs is set by: {caller_info.filename}")
24
+
25
+ if store_envs:
26
+ self.store_envs()
27
+
28
+ if secrets:
29
+ secrets_path = Path(__file__).parents[1] / "secrets.json"
30
+ if secrets_path.exists():
31
+ with open(secrets_path, "r") as rf:
32
+ secrets = json.load(rf)
33
+ else:
34
+ secrets = {}
35
+
36
+ if proxies:
37
+ for proxy_env in ["http_proxy", "https_proxy"]:
38
+ if isinstance(proxies, str):
39
+ self.envs[proxy_env] = proxies
40
+ elif "http_proxy" in secrets.keys():
41
+ self.envs[proxy_env] = secrets["http_proxy"]
42
+ elif os.getenv("http_proxy"):
43
+ self.envs[proxy_env] = os.getenv("http_proxy")
44
+ else:
45
+ continue
46
+
47
+ self.proxy = (
48
+ self.envs.get("all_proxy")
49
+ or self.envs.get("http_proxy")
50
+ or self.envs.get("https_proxy")
51
+ or None
52
+ )
53
+ self.requests_proxies = {
54
+ "http": self.proxy,
55
+ "https": self.proxy,
56
+ }
57
+
58
+ # https://www.proxynova.com/proxy-server-list/country-us/
59
+
60
+ print(f"Using proxy: [{self.proxy}]")
61
+ # r = requests.get(
62
+ # "http://ifconfig.me/ip",
63
+ # proxies=self.requests_proxies,
64
+ # timeout=10,
65
+ # )
66
+ # print(f"[r.status_code] r.text")
67
+
68
+
69
+ enver = OSEnver()
utils/enver.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import os
3
+
4
+ from pathlib import Path
5
+ from utils.logger import logger
6
+
7
+
8
+ class OSEnver:
9
+ def __init__(self):
10
+ self.envs_stack = []
11
+ self.envs = os.environ.copy()
12
+
13
+ def store_envs(self):
14
+ self.envs_stack.append(self.envs)
15
+
16
+ def restore_envs(self):
17
+ self.envs = self.envs_stack.pop()
18
+
19
+ def set_envs(self, secrets=True, proxies=None, store_envs=True):
20
+ # caller_info = inspect.stack()[1]
21
+ # logger.back(f"OS Envs is set by: {caller_info.filename}")
22
+
23
+ if store_envs:
24
+ self.store_envs()
25
+
26
+ if secrets:
27
+ secrets_path = Path(__file__).parents[1] / "secrets.json"
28
+ if secrets_path.exists():
29
+ with open(secrets_path, "r") as rf:
30
+ secrets = json.load(rf)
31
+ else:
32
+ secrets = {}
33
+
34
+ if proxies:
35
+ for proxy_env in ["http_proxy", "https_proxy"]:
36
+ if isinstance(proxies, str):
37
+ self.envs[proxy_env] = proxies
38
+ elif "http_proxy" in secrets.keys():
39
+ self.envs[proxy_env] = secrets["http_proxy"]
40
+ elif os.getenv("http_proxy"):
41
+ self.envs[proxy_env] = os.getenv("http_proxy")
42
+ else:
43
+ continue
44
+
45
+ self.proxy = (
46
+ self.envs.get("all_proxy")
47
+ or self.envs.get("http_proxy")
48
+ or self.envs.get("https_proxy")
49
+ or None
50
+ )
51
+ self.requests_proxies = {
52
+ "http": self.proxy,
53
+ "https": self.proxy,
54
+ }
55
+
56
+ if self.proxy:
57
+ logger.note(f"Using proxy: [{self.proxy}]")
58
+
59
+
60
+ enver = OSEnver()
utils/logger.py ADDED
@@ -0,0 +1,269 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import datetime
2
+ import functools
3
+ import inspect
4
+ import logging
5
+ import os
6
+ import shutil
7
+ import subprocess
8
+ from termcolor import colored
9
+
10
+
11
+ def add_fillers(text, filler="=", fill_side="both"):
12
+ terminal_width = shutil.get_terminal_size().columns
13
+ text = text.strip()
14
+ text_width = len(text)
15
+ if text_width >= terminal_width:
16
+ return text
17
+
18
+ if fill_side[0].lower() == "b":
19
+ leading_fill_str = filler * ((terminal_width - text_width) // 2 - 1) + " "
20
+ trailing_fill_str = " " + filler * (
21
+ terminal_width - text_width - len(leading_fill_str) - 1
22
+ )
23
+ elif fill_side[0].lower() == "l":
24
+ leading_fill_str = filler * (terminal_width - text_width - 1) + " "
25
+ trailing_fill_str = ""
26
+ elif fill_side[0].lower() == "r":
27
+ leading_fill_str = ""
28
+ trailing_fill_str = " " + filler * (terminal_width - text_width - 1)
29
+ else:
30
+ raise ValueError("Invalid fill_side")
31
+
32
+ filled_str = f"{leading_fill_str}{text}{trailing_fill_str}"
33
+ return filled_str
34
+
35
+
36
+ class OSLogger(logging.Logger):
37
+ LOG_METHODS = {
38
+ "err": ("error", "red"),
39
+ "warn": ("warning", "light_red"),
40
+ "note": ("info", "light_magenta"),
41
+ "mesg": ("info", "light_cyan"),
42
+ "file": ("info", "light_blue"),
43
+ "line": ("info", "white"),
44
+ "success": ("info", "light_green"),
45
+ "fail": ("info", "light_red"),
46
+ "back": ("debug", "light_cyan"),
47
+ }
48
+ INDENT_METHODS = [
49
+ "indent",
50
+ "set_indent",
51
+ "reset_indent",
52
+ "store_indent",
53
+ "restore_indent",
54
+ "log_indent",
55
+ ]
56
+ LEVEL_METHODS = [
57
+ "set_level",
58
+ "store_level",
59
+ "restore_level",
60
+ "quiet",
61
+ "enter_quiet",
62
+ "exit_quiet",
63
+ ]
64
+ LEVEL_NAMES = {
65
+ "critical": logging.CRITICAL,
66
+ "error": logging.ERROR,
67
+ "warning": logging.WARNING,
68
+ "info": logging.INFO,
69
+ "debug": logging.DEBUG,
70
+ }
71
+
72
+ def __init__(self, name=None, prefix=False):
73
+ if not name:
74
+ frame = inspect.stack()[1]
75
+ module = inspect.getmodule(frame[0])
76
+ name = module.__name__
77
+
78
+ super().__init__(name)
79
+ self.setLevel(logging.INFO)
80
+
81
+ if prefix:
82
+ formatter_prefix = "[%(asctime)s] - [%(name)s] - [%(levelname)s]\n"
83
+ else:
84
+ formatter_prefix = ""
85
+
86
+ self.formatter = logging.Formatter(formatter_prefix + "%(message)s")
87
+
88
+ stream_handler = logging.StreamHandler()
89
+ stream_handler.setLevel(logging.INFO)
90
+ stream_handler.setFormatter(self.formatter)
91
+ self.addHandler(stream_handler)
92
+
93
+ self.log_indent = 0
94
+ self.log_indents = []
95
+
96
+ self.log_level = "info"
97
+ self.log_levels = []
98
+
99
+ def indent(self, indent=2):
100
+ self.log_indent += indent
101
+
102
+ def set_indent(self, indent=2):
103
+ self.log_indent = indent
104
+
105
+ def reset_indent(self):
106
+ self.log_indent = 0
107
+
108
+ def store_indent(self):
109
+ self.log_indents.append(self.log_indent)
110
+
111
+ def restore_indent(self):
112
+ self.log_indent = self.log_indents.pop(-1)
113
+
114
+ def set_level(self, level):
115
+ self.log_level = level
116
+ self.setLevel(self.LEVEL_NAMES[level])
117
+
118
+ def store_level(self):
119
+ self.log_levels.append(self.log_level)
120
+
121
+ def restore_level(self):
122
+ self.log_level = self.log_levels.pop(-1)
123
+ self.set_level(self.log_level)
124
+
125
+ def quiet(self):
126
+ self.set_level("critical")
127
+
128
+ def enter_quiet(self, quiet=False):
129
+ if quiet:
130
+ self.store_level()
131
+ self.quiet()
132
+
133
+ def exit_quiet(self, quiet=False):
134
+ if quiet:
135
+ self.restore_level()
136
+
137
+ def log(
138
+ self,
139
+ level,
140
+ color,
141
+ msg,
142
+ indent=0,
143
+ fill=False,
144
+ fill_side="both",
145
+ end="\n",
146
+ *args,
147
+ **kwargs,
148
+ ):
149
+ if type(msg) == str:
150
+ msg_str = msg
151
+ else:
152
+ msg_str = repr(msg)
153
+ quotes = ["'", '"']
154
+ if msg_str[0] in quotes and msg_str[-1] in quotes:
155
+ msg_str = msg_str[1:-1]
156
+
157
+ indent_str = " " * (self.log_indent + indent)
158
+ indented_msg = "\n".join([indent_str + line for line in msg_str.split("\n")])
159
+
160
+ if fill:
161
+ indented_msg = add_fillers(indented_msg, fill_side=fill_side)
162
+
163
+ handler = self.handlers[0]
164
+ handler.terminator = end
165
+
166
+ getattr(self, level)(colored(indented_msg, color), *args, **kwargs)
167
+
168
+ def route_log(self, method, msg, *args, **kwargs):
169
+ level, method = method
170
+ functools.partial(self.log, level, method, msg)(*args, **kwargs)
171
+
172
+ def err(self, msg: str = "", *args, **kwargs):
173
+ self.route_log(("error", "red"), msg, *args, **kwargs)
174
+
175
+ def warn(self, msg: str = "", *args, **kwargs):
176
+ self.route_log(("warning", "light_red"), msg, *args, **kwargs)
177
+
178
+ def note(self, msg: str = "", *args, **kwargs):
179
+ self.route_log(("info", "light_magenta"), msg, *args, **kwargs)
180
+
181
+ def mesg(self, msg: str = "", *args, **kwargs):
182
+ self.route_log(("info", "light_cyan"), msg, *args, **kwargs)
183
+
184
+ def file(self, msg: str = "", *args, **kwargs):
185
+ self.route_log(("info", "light_blue"), msg, *args, **kwargs)
186
+
187
+ def line(self, msg: str = "", *args, **kwargs):
188
+ self.route_log(("info", "white"), msg, *args, **kwargs)
189
+
190
+ def success(self, msg: str = "", *args, **kwargs):
191
+ self.route_log(("info", "light_green"), msg, *args, **kwargs)
192
+
193
+ def fail(self, msg: str = "", *args, **kwargs):
194
+ self.route_log(("info", "light_red"), msg, *args, **kwargs)
195
+
196
+ def back(self, msg: str = "", *args, **kwargs):
197
+ self.route_log(("debug", "light_cyan"), msg, *args, **kwargs)
198
+
199
+
200
+ logger = OSLogger()
201
+
202
+
203
+ def shell_cmd(cmd, getoutput=False, showcmd=True, env=None):
204
+ if showcmd:
205
+ logger.info(colored(f"\n$ [{os.getcwd()}]", "light_blue"))
206
+ logger.info(colored(f" $ {cmd}\n", "light_cyan"))
207
+ if getoutput:
208
+ output = subprocess.getoutput(cmd, env=env)
209
+ return output
210
+ else:
211
+ subprocess.run(cmd, shell=True, env=env)
212
+
213
+
214
+ class Runtimer:
215
+ def __enter__(self):
216
+ self.t1, _ = self.start_time()
217
+ return self
218
+
219
+ def __exit__(self, exc_type, exc_value, traceback):
220
+ self.t2, _ = self.end_time()
221
+ self.elapsed_time(self.t2 - self.t1)
222
+
223
+ def start_time(self):
224
+ t1 = datetime.datetime.now()
225
+ self.logger_time("start", t1)
226
+ return t1, self.time2str(t1)
227
+
228
+ def end_time(self):
229
+ t2 = datetime.datetime.now()
230
+ self.logger_time("end", t2)
231
+ return t2, self.time2str(t2)
232
+
233
+ def elapsed_time(self, dt=None):
234
+ if dt is None:
235
+ dt = self.t2 - self.t1
236
+ self.logger_time("elapsed", dt)
237
+ return dt, self.time2str(dt)
238
+
239
+ def logger_time(self, time_type, t):
240
+ time_types = {
241
+ "start": "Start",
242
+ "end": "End",
243
+ "elapsed": "Elapsed",
244
+ }
245
+ time_str = add_fillers(
246
+ colored(
247
+ f"{time_types[time_type]} time: [ {self.time2str(t)} ]",
248
+ "light_magenta",
249
+ ),
250
+ fill_side="both",
251
+ )
252
+ logger.line(time_str)
253
+
254
+ # Convert time to string
255
+ def time2str(self, t):
256
+ datetime_str_format = "%Y-%m-%d %H:%M:%S"
257
+ if isinstance(t, datetime.datetime):
258
+ return t.strftime(datetime_str_format)
259
+ elif isinstance(t, datetime.timedelta):
260
+ hours = t.seconds // 3600
261
+ hour_str = f"{hours} hr" if hours > 0 else ""
262
+ minutes = (t.seconds // 60) % 60
263
+ minute_str = f"{minutes:>2} min" if minutes > 0 else ""
264
+ seconds = t.seconds % 60
265
+ second_str = f"{seconds:>2} s"
266
+ time_str = " ".join([hour_str, minute_str, second_str]).strip()
267
+ return time_str
268
+ else:
269
+ return str(t)