ka1kuk commited on
Commit
da27d7f
1 Parent(s): be19c13

Upload 21 files

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://Hansimov:$HF_TOKEN@huggingface.co/spaces/Hansimov/hf-llm-api 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,165 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: LLM API
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
+ ## Features
14
+
15
+ ✅ Implemented:
16
+
17
+ - Available Models:
18
+ - `mixtral-8x7b`, `mistral-7b`, `openchat-3.5`
19
+ - Adaptive prompt templates for different models
20
+ - Support OpenAI API format
21
+ - Can use 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
24
+ - Docker deployment
25
+
26
+ 🔨 In progress:
27
+ - [ ] Support more models
28
+ - [ ] meta-llama/Llama-2-70b-chat-hf
29
+ - [ ] codellama/CodeLlama-34b-Instruct-hf
30
+ - [ ] tiiuae/falcon-180B-chat
31
+
32
+
33
+ ## Run API service
34
+
35
+ ### Run in Command Line
36
+
37
+ **Install dependencies:**
38
+
39
+ ```bash
40
+ # pipreqs . --force --mode no-pin
41
+ pip install -r requirements.txt
42
+ ```
43
+
44
+ **Run API:**
45
+
46
+ ```bash
47
+ python -m apis.chat_api
48
+ ```
49
+
50
+ ## Run via Docker
51
+
52
+ **Docker build:**
53
+
54
+ ```bash
55
+ sudo docker build -t hf-llm-api:1.0 . --build-arg http_proxy=$http_proxy --build-arg https_proxy=$https_proxy
56
+ ```
57
+
58
+ **Docker run:**
59
+
60
+ ```bash
61
+ # no proxy
62
+ sudo docker run -p 23333:23333 hf-llm-api:1.0
63
+
64
+ # with proxy
65
+ sudo docker run -p 23333:23333 --env http_proxy="http://<server>:<port>" hf-llm-api:1.0
66
+ ```
67
+
68
+ ## API Usage
69
+
70
+ ### Using `openai-python`
71
+
72
+ See: [examples/chat_with_openai.py](https://github.com/Hansimov/hf-llm-api/blob/main/examples/chat_with_openai.py)
73
+
74
+ ```py
75
+ from openai import OpenAI
76
+
77
+ # If runnning this service with proxy, you might need to unset `http(s)_proxy`.
78
+ base_url = "http://127.0.0.1:23333"
79
+ # Your own HF_TOKEN
80
+ api_key = "hf_xxxxxxxxxxxxxxxx"
81
+
82
+ client = OpenAI(base_url=base_url, api_key=api_key)
83
+ response = client.chat.completions.create(
84
+ model="mixtral-8x7b",
85
+ messages=[
86
+ {
87
+ "role": "user",
88
+ "content": "what is your model",
89
+ }
90
+ ],
91
+ stream=True,
92
+ )
93
+
94
+ for chunk in response:
95
+ if chunk.choices[0].delta.content is not None:
96
+ print(chunk.choices[0].delta.content, end="", flush=True)
97
+ elif chunk.choices[0].finish_reason == "stop":
98
+ print()
99
+ else:
100
+ pass
101
+ ```
102
+
103
+ ### Using post requests
104
+
105
+ See: [examples/chat_with_post.py](https://github.com/Hansimov/hf-llm-api/blob/main/examples/chat_with_post.py)
106
+
107
+
108
+ ```py
109
+ import ast
110
+ import httpx
111
+ import json
112
+ import re
113
+
114
+ # If runnning this service with proxy, you might need to unset `http(s)_proxy`.
115
+ chat_api = "http://127.0.0.1:23333"
116
+ api_key = "sk-xxxxx"
117
+ requests_headers = {}
118
+ requests_payload = {
119
+ "model": "mixtral-8x7b",
120
+ "messages": [
121
+ {
122
+ "role": "user",
123
+ "content": "what is your model",
124
+ }
125
+ ],
126
+ "stream": True,
127
+ }
128
+
129
+ with httpx.stream(
130
+ "POST",
131
+ chat_api + "/chat/completions",
132
+ headers=requests_headers,
133
+ json=requests_payload,
134
+ timeout=httpx.Timeout(connect=20, read=60, write=20, pool=None),
135
+ ) as response:
136
+ # https://docs.aiohttp.org/en/stable/streams.html
137
+ # https://github.com/openai/openai-cookbook/blob/main/examples/How_to_stream_completions.ipynb
138
+ response_content = ""
139
+ for line in response.iter_lines():
140
+ remove_patterns = [r"^\s*data:\s*", r"^\s*\[DONE\]\s*"]
141
+ for pattern in remove_patterns:
142
+ line = re.sub(pattern, "", line).strip()
143
+
144
+ if line:
145
+ try:
146
+ line_data = json.loads(line)
147
+ except Exception as e:
148
+ try:
149
+ line_data = ast.literal_eval(line)
150
+ except:
151
+ print(f"Error: {line}")
152
+ raise e
153
+ # print(f"line: {line_data}")
154
+ delta_data = line_data["choices"][0]["delta"]
155
+ finish_reason = line_data["choices"][0]["finish_reason"]
156
+ if "role" in delta_data:
157
+ role = delta_data["role"]
158
+ if "content" in delta_data:
159
+ delta_content = delta_data["content"]
160
+ response_content += delta_content
161
+ print(delta_content, end="", flush=True)
162
+ if finish_reason == "stop":
163
+ print()
164
+
165
+ ```
__init__.py ADDED
File without changes
apis/__init__.py ADDED
File without changes
apis/chat_api.py ADDED
@@ -0,0 +1,155 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import os
3
+ import sys
4
+ import uvicorn
5
+
6
+ from fastapi import FastAPI, Depends
7
+ from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
8
+ from pydantic import BaseModel, Field
9
+ from sse_starlette.sse import EventSourceResponse, ServerSentEvent
10
+ from utils.logger import logger
11
+ from networks.message_streamer import MessageStreamer
12
+ from messagers.message_composer import MessageComposer
13
+ from mocks.stream_chat_mocker import stream_chat_mock
14
+
15
+
16
+ class ChatAPIApp:
17
+ def __init__(self):
18
+ self.app = FastAPI(
19
+ docs_url="/",
20
+ title="HuggingFace LLM API",
21
+ swagger_ui_parameters={"defaultModelsExpandDepth": -1},
22
+ version="1.0",
23
+ )
24
+ self.setup_routes()
25
+
26
+ def get_available_models(self):
27
+ # ANCHOR[id=available-models]: Available models
28
+ self.available_models = [
29
+ {
30
+ "id": "mixtral-8x7b",
31
+ "description": "[mistralai/Mixtral-8x7B-Instruct-v0.1]: https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1",
32
+ },
33
+ {
34
+ "id": "mistral-7b",
35
+ "description": "[mistralai/Mistral-7B-Instruct-v0.2]: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2",
36
+ },
37
+ {
38
+ "id": "openchat-3.5",
39
+ "description": "[openchat/openchat-3.5-1210]: https://huggingface.co/openchat/openchat-3.5-1210",
40
+ },
41
+ ]
42
+ return self.available_models
43
+
44
+ def extract_api_key(
45
+ credentials: HTTPAuthorizationCredentials = Depends(
46
+ HTTPBearer(auto_error=False)
47
+ ),
48
+ ):
49
+ if credentials:
50
+ return credentials.credentials
51
+ else:
52
+ return os.getenv("HF_TOKEN") or None
53
+
54
+ class ChatCompletionsPostItem(BaseModel):
55
+ model: str = Field(
56
+ default="mixtral-8x7b",
57
+ description="(str) `mixtral-8x7b`",
58
+ )
59
+ messages: list = Field(
60
+ default=[{"role": "user", "content": "Hello, who are you?"}],
61
+ description="(list) Messages",
62
+ )
63
+ temperature: float = Field(
64
+ default=0.01,
65
+ description="(float) Temperature",
66
+ )
67
+ max_tokens: int = Field(
68
+ default=4096,
69
+ description="(int) Max tokens",
70
+ )
71
+ stream: bool = Field(
72
+ default=True,
73
+ description="(bool) Stream",
74
+ )
75
+
76
+ def chat_completions(
77
+ self, item: ChatCompletionsPostItem, api_key: str = Depends(extract_api_key)
78
+ ):
79
+ streamer = MessageStreamer(model=item.model)
80
+ composer = MessageComposer(model=item.model)
81
+ composer.merge(messages=item.messages)
82
+ # streamer.chat = stream_chat_mock
83
+
84
+ stream_response = streamer.chat_response(
85
+ prompt=composer.merged_str,
86
+ temperature=item.temperature,
87
+ max_new_tokens=item.max_tokens,
88
+ api_key=api_key,
89
+ )
90
+ if item.stream:
91
+ event_source_response = EventSourceResponse(
92
+ streamer.chat_return_generator(stream_response),
93
+ media_type="text/event-stream",
94
+ ping=2000,
95
+ ping_message_factory=lambda: ServerSentEvent(**{"comment": ""}),
96
+ )
97
+ return event_source_response
98
+ else:
99
+ data_response = streamer.chat_return_dict(stream_response)
100
+ return data_response
101
+
102
+ def setup_routes(self):
103
+ for prefix in ["", "/v1"]:
104
+ self.app.get(
105
+ prefix + "/models",
106
+ summary="Get available models",
107
+ )(self.get_available_models)
108
+
109
+ self.app.post(
110
+ prefix + "/chat/completions",
111
+ summary="Chat completions in conversation session",
112
+ )(self.chat_completions)
113
+
114
+
115
+ class ArgParser(argparse.ArgumentParser):
116
+ def __init__(self, *args, **kwargs):
117
+ super(ArgParser, self).__init__(*args, **kwargs)
118
+
119
+ self.add_argument(
120
+ "-s",
121
+ "--server",
122
+ type=str,
123
+ default="0.0.0.0",
124
+ help="Server IP for HF LLM Chat API",
125
+ )
126
+ self.add_argument(
127
+ "-p",
128
+ "--port",
129
+ type=int,
130
+ default=23333,
131
+ help="Server Port for HF LLM Chat API",
132
+ )
133
+
134
+ self.add_argument(
135
+ "-d",
136
+ "--dev",
137
+ default=False,
138
+ action="store_true",
139
+ help="Run in dev mode",
140
+ )
141
+
142
+ self.args = self.parse_args(sys.argv[1:])
143
+
144
+
145
+ app = ChatAPIApp().app
146
+
147
+ if __name__ == "__main__":
148
+ args = ArgParser().args
149
+ if args.dev:
150
+ uvicorn.run("__main__:app", host=args.server, port=args.port, reload=True)
151
+ else:
152
+ uvicorn.run("__main__:app", host=args.server, port=args.port, reload=False)
153
+
154
+ # python -m apis.chat_api # [Docker] on product mode
155
+ # 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,203 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ from pprint import pprint
3
+
4
+
5
+ class MessageComposer:
6
+ # LINK - apis/chat_api.py#available-models
7
+ AVALAIBLE_MODELS = [
8
+ "mixtral-8x7b",
9
+ "mistral-7b",
10
+ "openchat-3.5",
11
+ ]
12
+
13
+ def __init__(self, model: str = None):
14
+ if model in self.AVALAIBLE_MODELS:
15
+ self.model = model
16
+ else:
17
+ self.model = "mixtral-8x7b"
18
+ self.inst_roles = ["user", "system", "inst"]
19
+ self.answer_roles = ["assistant", "bot", "answer"]
20
+
21
+ def concat_messages_by_role(self, messages):
22
+ def is_same_role(role1, role2):
23
+ if (
24
+ (role1 == role2)
25
+ or (role1 in self.inst_roles and role2 in self.inst_roles)
26
+ or (role1 in self.answer_roles and role2 in self.answer_roles)
27
+ ):
28
+ return True
29
+ else:
30
+ return False
31
+
32
+ concat_messages = []
33
+ for message in messages:
34
+ role = message["role"]
35
+ content = message["content"]
36
+ if concat_messages and is_same_role(role, concat_messages[-1]["role"]):
37
+ concat_messages[-1]["content"] += "\n" + content
38
+ else:
39
+ if role in self.inst_roles:
40
+ message["role"] = "inst"
41
+ elif role in self.answer_roles:
42
+ message["role"] = "answer"
43
+ else:
44
+ message["role"] = "inst"
45
+ concat_messages.append(message)
46
+ return concat_messages
47
+
48
+ def merge(self, messages) -> str:
49
+ # Mistral and Mixtral:
50
+ # <s> [INST] Instruction [/INST] Model answer </s> [INST] Follow-up instruction [/INST]
51
+ # OpenChat:
52
+ # 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:
53
+
54
+ self.messages = self.concat_messages_by_role(messages)
55
+ self.merged_str = ""
56
+
57
+ if self.model in ["mixtral-8x7b", "mistral-7b"]:
58
+ self.cached_str = ""
59
+ for message in self.messages:
60
+ role = message["role"]
61
+ content = message["content"]
62
+ if role in self.inst_roles:
63
+ self.cached_str = f"[INST] {content} [/INST]"
64
+ elif role in self.answer_roles:
65
+ self.merged_str += f"<s> {self.cached_str} {content} </s>\n"
66
+ self.cached_str = ""
67
+ else:
68
+ self.cached_str = f"[INST] {content} [/INST]"
69
+ if self.cached_str:
70
+ self.merged_str += f"{self.cached_str}"
71
+ elif self.model in ["openchat-3.5"]:
72
+ self.merged_str_list = []
73
+ self.end_of_turn = "<|end_of_turn|>"
74
+ for message in self.messages:
75
+ role = message["role"]
76
+ content = message["content"]
77
+ if role in self.inst_roles:
78
+ self.merged_str_list.append(
79
+ f"GPT4 Correct User:\n{content}{self.end_of_turn}"
80
+ )
81
+ elif role in self.answer_roles:
82
+ self.merged_str_list.append(
83
+ f"GPT4 Correct Assistant:\n{content}{self.end_of_turn}"
84
+ )
85
+ else:
86
+ self.merged_str_list.append(
87
+ f"GPT4 Correct User: {content}{self.end_of_turn}"
88
+ )
89
+ self.merged_str_list.append(f"GPT4 Correct Assistant:\n")
90
+ self.merged_str = "\n".join(self.merged_str_list)
91
+ else:
92
+ self.merged_str = "\n".join(
93
+ [
94
+ f'`{message["role"]}`:\n{message["content"]}\n'
95
+ for message in self.messages
96
+ ]
97
+ )
98
+
99
+ return self.merged_str
100
+
101
+ def convert_pair_matches_to_messages(self, pair_matches_list):
102
+ messages = []
103
+ if len(pair_matches_list) <= 0:
104
+ messages = [
105
+ {
106
+ "role": "user",
107
+ "content": self.merged_str,
108
+ }
109
+ ]
110
+ else:
111
+ for match in pair_matches_list:
112
+ inst = match.group("inst")
113
+ answer = match.group("answer")
114
+ messages.extend(
115
+ [
116
+ {"role": "user", "content": inst.strip()},
117
+ {"role": "assistant", "content": answer.strip()},
118
+ ]
119
+ )
120
+ return messages
121
+
122
+ def append_last_instruction_to_messages(self, inst_matches_list, pair_matches_list):
123
+ if len(inst_matches_list) > len(pair_matches_list):
124
+ self.messages.extend(
125
+ [
126
+ {
127
+ "role": "user",
128
+ "content": inst_matches_list[-1].group("inst").strip(),
129
+ }
130
+ ]
131
+ )
132
+
133
+ def split(self, merged_str) -> list:
134
+ self.merged_str = merged_str
135
+ self.messages = []
136
+
137
+ if self.model in ["mixtral-8x7b", "mistral-7b"]:
138
+ pair_pattern = (
139
+ r"<s>\s*\[INST\](?P<inst>[\s\S]*?)\[/INST\](?P<answer>[\s\S]*?)</s>"
140
+ )
141
+ pair_matches = re.finditer(pair_pattern, self.merged_str, re.MULTILINE)
142
+ pair_matches_list = list(pair_matches)
143
+
144
+ self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
145
+
146
+ inst_pattern = r"\[INST\](?P<inst>[\s\S]*?)\[/INST\]"
147
+ inst_matches = re.finditer(inst_pattern, self.merged_str, re.MULTILINE)
148
+ inst_matches_list = list(inst_matches)
149
+
150
+ self.append_last_instruction_to_messages(
151
+ inst_matches_list, pair_matches_list
152
+ )
153
+
154
+ elif self.model in ["openchat-3.5"]:
155
+ pair_pattern = r"GPT4 Correct User:(?P<inst>[\s\S]*?)<\|end_of_turn\|>\s*GPT4 Correct Assistant:(?P<answer>[\s\S]*?)<\|end_of_turn\|>"
156
+ # ignore case
157
+ pair_matches = re.finditer(
158
+ pair_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
159
+ )
160
+ pair_matches_list = list(pair_matches)
161
+ self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
162
+ inst_pattern = r"GPT4 Correct User:(?P<inst>[\s\S]*?)<\|end_of_turn\|>"
163
+ inst_matches = re.finditer(
164
+ inst_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
165
+ )
166
+ inst_matches_list = list(inst_matches)
167
+ self.append_last_instruction_to_messages(
168
+ inst_matches_list, pair_matches_list
169
+ )
170
+ else:
171
+ self.messages = [
172
+ {
173
+ "role": "user",
174
+ "content": self.merged_str,
175
+ }
176
+ ]
177
+
178
+ return self.messages
179
+
180
+
181
+ if __name__ == "__main__":
182
+ composer = MessageComposer(model="openchat-3.5")
183
+ messages = [
184
+ {
185
+ "role": "system",
186
+ "content": "You are a LLM developed by OpenAI. Your name is GPT-4.",
187
+ },
188
+ {"role": "user", "content": "Hello, who are you?"},
189
+ {"role": "assistant", "content": "I am a bot."},
190
+ {"role": "user", "content": "What is your name?"},
191
+ # {"role": "assistant", "content": "My name is Bing."},
192
+ # {"role": "user", "content": "Tell me a joke."},
193
+ # {"role": "assistant", "content": "What is a robot's favorite type of music?"},
194
+ # {
195
+ # "role": "user",
196
+ # "content": "How many questions have I asked? Please list them.",
197
+ # },
198
+ ]
199
+ print("model:", composer.model)
200
+ merged_str = composer.merge(messages)
201
+ print(merged_str)
202
+ pprint(composer.split(merged_str))
203
+ # print(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,156 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import re
3
+ import requests
4
+ from messagers.message_outputer import OpenaiStreamOutputer
5
+ from utils.logger import logger
6
+ from utils.enver import enver
7
+
8
+
9
+ class MessageStreamer:
10
+ MODEL_MAP = {
11
+ "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1", # 72.62, fast [Recommended]
12
+ "mistral-7b": "mistralai/Mistral-7B-Instruct-v0.2", # 65.71, fast
13
+ "openchat-3.5": "openchat/openchat-3.5-1210", # ??, fast
14
+ # "zephyr-7b-alpha": "HuggingFaceH4/zephyr-7b-alpha", # 59.5, fast
15
+ # "zephyr-7b-beta": "HuggingFaceH4/zephyr-7b-beta", # 61.95, slow
16
+ "default": "mistralai/Mixtral-8x7B-Instruct-v0.1",
17
+ }
18
+ STOP_SEQUENCES_MAP = {
19
+ "mixtral-8x7b": "</s>",
20
+ "mistral-7b": "</s>",
21
+ "openchat-3.5": "<|end_of_turn|>",
22
+ }
23
+
24
+ def __init__(self, model: str):
25
+ if model in self.MODEL_MAP.keys():
26
+ self.model = model
27
+ else:
28
+ self.model = "default"
29
+ self.model_fullname = self.MODEL_MAP[self.model]
30
+ self.message_outputer = OpenaiStreamOutputer()
31
+
32
+ def parse_line(self, line):
33
+ line = line.decode("utf-8")
34
+ line = re.sub(r"data:\s*", "", line)
35
+ data = json.loads(line)
36
+ content = data["token"]["text"]
37
+ return content
38
+
39
+ def chat_response(
40
+ self,
41
+ prompt: str = None,
42
+ temperature: float = 0.01,
43
+ max_new_tokens: int = 8192,
44
+ api_key: str = None,
45
+ ):
46
+ # https://huggingface.co/docs/api-inference/detailed_parameters?code=curl
47
+ # 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>'
48
+ self.request_url = (
49
+ f"https://api-inference.huggingface.co/models/{self.model_fullname}"
50
+ )
51
+ self.request_headers = {
52
+ "Content-Type": "application/json",
53
+ }
54
+
55
+ if api_key:
56
+ logger.note(
57
+ f"Using API Key: {api_key[:3]}{(len(api_key)-7)*'*'}{api_key[-4:]}"
58
+ )
59
+ self.request_headers["Authorization"] = f"Bearer {api_key}"
60
+
61
+ # References:
62
+ # huggingface_hub/inference/_client.py:
63
+ # class InferenceClient > def text_generation()
64
+ # huggingface_hub/inference/_text_generation.py:
65
+ # class TextGenerationRequest > param `stream`
66
+ # https://huggingface.co/docs/text-generation-inference/conceptual/streaming#streaming-with-curl
67
+ self.request_body = {
68
+ "inputs": prompt,
69
+ "parameters": {
70
+ "temperature": max(temperature, 0.01), # must be positive
71
+ "max_new_tokens": max_new_tokens,
72
+ "return_full_text": False,
73
+ },
74
+ "stream": True,
75
+ }
76
+
77
+ if self.model in self.STOP_SEQUENCES_MAP.keys():
78
+ self.stop_sequences = self.STOP_SEQUENCES_MAP[self.model]
79
+ # self.request_body["parameters"]["stop_sequences"] = [
80
+ # self.STOP_SEQUENCES[self.model]
81
+ # ]
82
+
83
+ logger.back(self.request_url)
84
+ enver.set_envs(proxies=True)
85
+ stream_response = requests.post(
86
+ self.request_url,
87
+ headers=self.request_headers,
88
+ json=self.request_body,
89
+ proxies=enver.requests_proxies,
90
+ stream=True,
91
+ )
92
+ status_code = stream_response.status_code
93
+ if status_code == 200:
94
+ logger.success(status_code)
95
+ else:
96
+ logger.err(status_code)
97
+
98
+ return stream_response
99
+
100
+ def chat_return_dict(self, stream_response):
101
+ # https://platform.openai.com/docs/guides/text-generation/chat-completions-response-format
102
+ final_output = self.message_outputer.default_data.copy()
103
+ final_output["choices"] = [
104
+ {
105
+ "index": 0,
106
+ "finish_reason": "stop",
107
+ "message": {
108
+ "role": "assistant",
109
+ "content": "",
110
+ },
111
+ }
112
+ ]
113
+ logger.back(final_output)
114
+
115
+ final_content = ""
116
+ for line in stream_response.iter_lines():
117
+ if not line:
118
+ continue
119
+ content = self.parse_line(line)
120
+
121
+ if content.strip() == self.stop_sequences:
122
+ logger.success("\n[Finished]")
123
+ break
124
+ else:
125
+ logger.back(content, end="")
126
+ final_content += content
127
+
128
+ if self.model in self.STOP_SEQUENCES_MAP.keys():
129
+ final_content = final_content.replace(self.stop_sequences, "")
130
+
131
+ final_output["choices"][0]["message"]["content"] = final_content
132
+ return final_output
133
+
134
+ def chat_return_generator(self, stream_response):
135
+ is_finished = False
136
+ for line in stream_response.iter_lines():
137
+ if not line:
138
+ continue
139
+
140
+ content = self.parse_line(line)
141
+
142
+ if content.strip() == self.stop_sequences:
143
+ content_type = "Finished"
144
+ logger.success("\n[Finished]")
145
+ is_finished = True
146
+ else:
147
+ content_type = "Completions"
148
+ logger.back(content, end="")
149
+
150
+ output = self.message_outputer.output(
151
+ content=content, content_type=content_type
152
+ )
153
+ yield output
154
+
155
+ if not is_finished:
156
+ yield self.message_outputer.output(content="", content_type="Finished")
requirements.txt ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ aiohttp
2
+ fastapi
3
+ httpx
4
+ openai
5
+ pydantic
6
+ requests
7
+ sse_starlette
8
+ termcolor
9
+ uvicorn
10
+ 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)