Yarik commited on
Commit
de9ce4c
·
0 Parent(s):

Initial commit

Browse files
.gitattributes ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
.gitignore ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ secrets.json
2
+ .env
3
+ __pycache__
4
+ post.py
5
+ post_old.py
.idea/.gitignore ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ # Default ignored files
2
+ /shelf/
3
+ /workspace.xml
4
+ # Editor-based HTTP Client requests
5
+ /httpRequests/
6
+ # Datasource local storage ignored files
7
+ /dataSources/
8
+ /dataSources.local.xml
.idea/chatGPT.iml ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ <?xml version="1.0" encoding="UTF-8"?>
2
+ <module type="PYTHON_MODULE" version="4">
3
+ <component name="NewModuleRootManager">
4
+ <content url="file://$MODULE_DIR$" />
5
+ <orderEntry type="inheritedJdk" />
6
+ <orderEntry type="sourceFolder" forTests="false" />
7
+ </component>
8
+ </module>
.idea/inspectionProfiles/Project_Default.xml ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ <component name="InspectionProjectProfileManager">
2
+ <profile version="1.0">
3
+ <option name="myName" value="Project Default" />
4
+ <inspection_tool class="Eslint" enabled="true" level="WARNING" enabled_by_default="true" />
5
+ <inspection_tool class="PyInterpreterInspection" enabled="false" level="WARNING" enabled_by_default="false" />
6
+ <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
7
+ <option name="ignoredPackages">
8
+ <value>
9
+ <list size="13">
10
+ <item index="0" class="java.lang.String" itemvalue="httpx" />
11
+ <item index="1" class="java.lang.String" itemvalue="tiktoken" />
12
+ <item index="2" class="java.lang.String" itemvalue="pydantic" />
13
+ <item index="3" class="java.lang.String" itemvalue="markdown2" />
14
+ <item index="4" class="java.lang.String" itemvalue="requests" />
15
+ <item index="5" class="java.lang.String" itemvalue="openai" />
16
+ <item index="6" class="java.lang.String" itemvalue="sse_starlette" />
17
+ <item index="7" class="java.lang.String" itemvalue="termcolor" />
18
+ <item index="8" class="java.lang.String" itemvalue="fastapi" />
19
+ <item index="9" class="java.lang.String" itemvalue="aiohttp" />
20
+ <item index="10" class="java.lang.String" itemvalue="uvicorn" />
21
+ <item index="11" class="java.lang.String" itemvalue="websockets" />
22
+ <item index="12" class="java.lang.String" itemvalue="passlib" />
23
+ </list>
24
+ </value>
25
+ </option>
26
+ </inspection_tool>
27
+ </profile>
28
+ </component>
.idea/inspectionProfiles/profiles_settings.xml ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ <component name="InspectionProjectProfileManager">
2
+ <settings>
3
+ <option name="USE_PROJECT_PROFILE" value="false" />
4
+ <version value="1.0" />
5
+ </settings>
6
+ </component>
.idea/misc.xml ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ <?xml version="1.0" encoding="UTF-8"?>
2
+ <project version="4">
3
+ <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.12" project-jdk-type="Python SDK" />
4
+ </project>
.idea/modules.xml ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ <?xml version="1.0" encoding="UTF-8"?>
2
+ <project version="4">
3
+ <component name="ProjectModuleManager">
4
+ <modules>
5
+ <module fileurl="file://$PROJECT_DIR$/.idea/chatGPT.iml" filepath="$PROJECT_DIR$/.idea/chatGPT.iml" />
6
+ </modules>
7
+ </component>
8
+ </project>
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 7860
7
+ CMD ["python", "-m", "apis.chat_api"]
README.md ADDED
@@ -0,0 +1,165 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: Xche Ai
3
+ emoji: 👁
4
+ colorFrom: pink
5
+ colorTo: indigo
6
+ sdk: docker
7
+ pinned: false
8
+ license: apache-2.0
9
+ ---
10
+
11
+ ## HF-LLM-API
12
+ Huggingface LLM Inference API in OpenAI message format.
13
+
14
+
15
+
16
+ ## Features
17
+
18
+ - Available Models (2024/01/22):
19
+ - `mistral-7b`, `mixtral-8x7b`, `nous-mixtral-8x7b`
20
+ - Adaptive prompt templates for different models
21
+ - Support OpenAI API format
22
+ - Enable api endpoint via official `openai-python` package
23
+ - Support both stream and no-stream response
24
+ - Support API Key via both HTTP auth header and env varible
25
+ - Docker deployment
26
+
27
+ ## Run API service
28
+
29
+ ### Run in Command Line
30
+
31
+ **Install dependencies:**
32
+
33
+ ```bash
34
+ # pipreqs . --force --mode no-pin
35
+ pip install -r requirements.txt
36
+ ```
37
+
38
+ **Run API:**
39
+
40
+ ```bash
41
+ python -m apis.chat_api
42
+ ```
43
+
44
+ ## Run via Docker
45
+
46
+ **Docker build:**
47
+
48
+ ```bash
49
+ sudo docker build -t hf-llm-api:1.0 . --build-arg http_proxy=$http_proxy --build-arg https_proxy=$https_proxy
50
+ ```
51
+
52
+ **Docker run:**
53
+
54
+ ```bash
55
+ # no proxy
56
+ sudo docker run -p 23333:23333 hf-llm-api:1.0
57
+
58
+ # with proxy
59
+ sudo docker run -p 23333:23333 --env http_proxy="http://<server>:<port>" hf-llm-api:1.0
60
+ ```
61
+
62
+ ## API Usage
63
+
64
+ ### Using `openai-python`
65
+
66
+ See: [`examples/chat_with_openai.py`](https://github.com/ruslanmv/hf-llm-api-collection/blob/main/examples/chat_with_openai.py)
67
+
68
+ ```py
69
+ from openai import OpenAI
70
+
71
+ # If runnning this service with proxy, you might need to unset `http(s)_proxy`.
72
+ base_url = "http://127.0.0.1:23333"
73
+ # Your own HF_TOKEN
74
+ api_key = "hf_xxxxxxxxxxxxxxxx"
75
+ # use below as non-auth user
76
+ # api_key = "sk-xxx"
77
+
78
+ client = OpenAI(base_url=base_url, api_key=api_key)
79
+ response = client.chat.completions.create(
80
+ model="mixtral-8x7b",
81
+ messages=[
82
+ {
83
+ "role": "user",
84
+ "content": "what is your model",
85
+ }
86
+ ],
87
+ stream=True,
88
+ )
89
+
90
+ for chunk in response:
91
+ if chunk.choices[0].delta.content is not None:
92
+ print(chunk.choices[0].delta.content, end="", flush=True)
93
+ elif chunk.choices[0].finish_reason == "stop":
94
+ print()
95
+ else:
96
+ pass
97
+ ```
98
+
99
+ ### Using post requests
100
+
101
+ See: [`examples/chat_with_post.py`](https://github.com/ruslanmv/hf-llm-api-collection/blob/main/examples/chat_with_post.py)
102
+
103
+
104
+ ```py
105
+ import ast
106
+ import httpx
107
+ import json
108
+ import re
109
+
110
+ # If runnning this service with proxy, you might need to unset `http(s)_proxy`.
111
+ chat_api = "http://127.0.0.1:23333"
112
+ # Your own HF_TOKEN
113
+ api_key = "hf_xxxxxxxxxxxxxxxx"
114
+ # use below as non-auth user
115
+ # api_key = "sk-xxx"
116
+
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/.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ .env
2
+ chat_api_old.py
apis/__init__.py ADDED
File without changes
apis/chat_api.py ADDED
@@ -0,0 +1,203 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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, HTTPException
9
+ from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
10
+ from fastapi.responses import HTMLResponse
11
+ from pydantic import BaseModel, Field
12
+ from typing import Union
13
+ from sse_starlette.sse import EventSourceResponse, ServerSentEvent
14
+ from passlib.context import CryptContext
15
+ from utils.logger import logger
16
+ from networks.message_streamer import MessageStreamer
17
+ from messagers.message_composer import MessageComposer
18
+ from mocks.stream_chat_mocker import stream_chat_mock
19
+
20
+
21
+ # Create FastAPI app
22
+ app = FastAPI(docs_url=None, redoc_url=None)
23
+
24
+ oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
25
+
26
+ class Auth(BaseModel):
27
+ api_key: str
28
+ password: str
29
+
30
+ # Password hashing context
31
+ pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
32
+
33
+ def get_password_hash(password):
34
+ return pwd_context.hash(password)
35
+
36
+ def verify_password(plain_password, hashed_password):
37
+ return pwd_context.verify(plain_password, hashed_password)
38
+
39
+ # Fetch API key and password from environment variables
40
+ api_key = os.getenv("XCHE_API_KEY")
41
+ password = os.getenv("XCHE_PASSWORD")
42
+
43
+ # Pre-hash the password and store it in a fake database
44
+ fake_data_db = {
45
+ api_key: {
46
+ "api_key": api_key,
47
+ "password": get_password_hash(password) # Pre-hashed password
48
+ }
49
+ }
50
+
51
+
52
+
53
+ def get_api_key(db, api_key: str):
54
+ if api_key in db:
55
+ api_dict = db[api_key]
56
+ return Auth(**api_dict)
57
+
58
+ def authenticate(fake_db, api_key: str, password: str):
59
+ api_data = get_api_key(fake_db, api_key)
60
+ if not api_data:
61
+ return False
62
+ if not verify_password(password, api_data.password):
63
+ return False
64
+ return api_data
65
+
66
+
67
+
68
+
69
+ @app.post("/token")
70
+ async def login(form_data: OAuth2PasswordRequestForm = Depends()):
71
+ api_data = authenticate(fake_data_db, form_data.username, form_data.password)
72
+ if not api_data:
73
+ raise HTTPException(
74
+ status_code=400,
75
+ detail="Incorrect API KEY or Password",
76
+ headers={"WWW-Authenticate": "Bearer"},
77
+ )
78
+ return {"access_token": api_data.api_key, "token_type": "bearer"}
79
+
80
+ def check_api_token(token: str = Depends(oauth2_scheme)):
81
+ api_data = get_api_key(fake_data_db, token)
82
+ if not api_data:
83
+ raise HTTPException(status_code=403, detail="Invalid or missing API Key")
84
+ return api_data
85
+
86
+
87
+ class ChatAPIApp:
88
+ def __init__(self):
89
+ self.app = app
90
+ self.setup_routes()
91
+
92
+ # "mixtral-8x7b"
93
+ # "mistral - 7b"
94
+ # "nous-mixtral-8x7b"
95
+ # "zephyr-7b-beta"
96
+ # "starchat2-15b-v0.1"
97
+
98
+ class ChatCompletionsPostItem(BaseModel):
99
+ model: str = Field(
100
+ default="mixtral-8x7b",
101
+ description="(str) `mixtral-8x7b`",
102
+ )
103
+ messages: list = Field(
104
+ default=[{"role": "user", "content": "Hello, who are you?"}],
105
+ description="(list) Messages",
106
+ )
107
+ temperature: Union[float, None] = Field(
108
+ default=0.5,
109
+ description="(float) Temperature",
110
+ )
111
+ top_p: Union[float, None] = Field(
112
+ default=0.95,
113
+ description="(float) top p",
114
+ )
115
+ max_tokens: Union[int, None] = Field(
116
+ default=-1,
117
+ description="(int) Max tokens",
118
+ )
119
+ use_cache: bool = Field(
120
+ default=False,
121
+ description="(bool) Use cache",
122
+ )
123
+ stream: bool = Field(
124
+ default=True,
125
+ description="(bool) Stream",
126
+ )
127
+
128
+ def chat_completions(self, item: ChatCompletionsPostItem):
129
+ streamer = MessageStreamer(model=item.model)
130
+ composer = MessageComposer(model=item.model)
131
+ composer.merge(messages=item.messages)
132
+ # streamer.chat = stream_chat_mock
133
+
134
+ stream_response = streamer.chat_response(
135
+ prompt=composer.merged_str,
136
+ temperature=item.temperature,
137
+ top_p=item.top_p,
138
+ max_new_tokens=item.max_tokens,
139
+ use_cache=item.use_cache,
140
+ )
141
+ if item.stream:
142
+ event_source_response = EventSourceResponse(
143
+ streamer.chat_return_generator(stream_response),
144
+ media_type="text/event-stream",
145
+ ping=2000,
146
+ ping_message_factory=lambda: ServerSentEvent(**{"comment": ""}),
147
+ )
148
+ return event_source_response
149
+ else:
150
+ data_response = streamer.chat_return_dict(stream_response)
151
+ return data_response
152
+
153
+
154
+ def setup_routes(self):
155
+ for prefix in ["", "/v1", "/api", "/api/v1"]:
156
+ if prefix == "/api/v1":
157
+ include_in_schema = True
158
+ else:
159
+ include_in_schema = False
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
+
167
+ class ArgParser(argparse.ArgumentParser):
168
+ def __init__(self, *args, **kwargs):
169
+ super(ArgParser, self).__init__(*args, **kwargs)
170
+
171
+ self.add_argument(
172
+ "-s",
173
+ "--server",
174
+ type=str,
175
+ default="0.0.0.0",
176
+ help="Server IP for HF LLM Chat API",
177
+ )
178
+ self.add_argument(
179
+ "-p",
180
+ "--port",
181
+ type=int,
182
+ default=7860,
183
+ help="Server Port for HF LLM Chat API",
184
+ )
185
+
186
+ self.add_argument(
187
+ "-d",
188
+ "--dev",
189
+ default=False,
190
+ action="store_true",
191
+ help="Run in dev mode",
192
+ )
193
+
194
+ self.args = self.parse_args(sys.argv[1:])
195
+
196
+ app = ChatAPIApp().app
197
+
198
+ if __name__ == "__main__":
199
+ args = ArgParser().args
200
+ if args.dev:
201
+ uvicorn.run("__main__:app", host=args.server, port=args.port, reload=True)
202
+ else:
203
+ uvicorn.run("__main__:app", host=args.server, port=args.port, reload=False)
main.py ADDED
File without changes
messagers/__init__.py ADDED
File without changes
messagers/message_composer.py ADDED
@@ -0,0 +1,265 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ "zephyr-7b-beta",
12
+ "starchat2-15b-v0.1",
13
+ "nous-mixtral-8x7b",
14
+ ]
15
+
16
+ def __init__(self, model: str = None):
17
+ if model in self.AVALAIBLE_MODELS:
18
+ self.model = model
19
+ else:
20
+ self.model = "mixtral-8x7b"
21
+ self.system_roles = ["system"]
22
+ self.inst_roles = ["user", "system", "inst"]
23
+ self.answer_roles = ["assistant", "bot", "answer"]
24
+ self.default_role = "user"
25
+
26
+ def concat_messages_by_role(self, messages):
27
+ def is_same_role(role1, role2):
28
+ if (
29
+ (role1 == role2)
30
+ or (role1 in self.inst_roles and role2 in self.inst_roles)
31
+ or (role1 in self.answer_roles and role2 in self.answer_roles)
32
+ ):
33
+ return True
34
+ else:
35
+ return False
36
+
37
+ concat_messages = []
38
+ for message in messages:
39
+ role = message["role"]
40
+ content = message["content"]
41
+ if concat_messages and is_same_role(role, concat_messages[-1]["role"]):
42
+ concat_messages[-1]["content"] += "\n" + content
43
+ else:
44
+ if role in self.inst_roles:
45
+ message["role"] = "inst"
46
+ elif role in self.answer_roles:
47
+ message["role"] = "answer"
48
+ else:
49
+ message["role"] = "inst"
50
+ concat_messages.append(message)
51
+ return concat_messages
52
+
53
+ def merge(self, messages) -> str:
54
+ # Mistral and Mixtral:
55
+ # <s> [INST] Instruction [/INST] Model answer </s> [INST] Follow-up instruction [/INST]
56
+
57
+ # OpenChat:
58
+ # 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:
59
+
60
+ # Nous Mixtral:
61
+ # <|im_start|>system
62
+ # You are "Hermes 2".<|im_end|>
63
+ # <|im_start|>user
64
+ # Hello, who are you?<|im_end|>
65
+ # <|im_start|>assistant
66
+
67
+ self.messages = messages
68
+ self.merged_str = ""
69
+
70
+ # https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1#instruction-format
71
+ if self.model in ["mixtral-8x7b", "mistral-7b"]:
72
+ self.messages = self.concat_messages_by_role(messages)
73
+ self.cached_str = ""
74
+ for message in self.messages:
75
+ role = message["role"]
76
+ content = message["content"]
77
+ if role in self.inst_roles:
78
+ self.cached_str = f"[INST] {content} [/INST]"
79
+ elif role in self.answer_roles:
80
+ self.merged_str += f"<s> {self.cached_str} {content} </s>\n"
81
+ self.cached_str = ""
82
+ else:
83
+ self.cached_str = f"[INST] {content} [/INST]"
84
+ if self.cached_str:
85
+ self.merged_str += f"{self.cached_str}"
86
+ # https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO#prompt-format
87
+ elif self.model in ["nous-mixtral-8x7b"]:
88
+ self.merged_str_list = []
89
+ for message in self.messages:
90
+ role = message["role"]
91
+ content = message["content"]
92
+ if role not in ["system", "user", "assistant"]:
93
+ role = self.default_role
94
+ message_line = f"<|im_start|>{role}\n{content}<|im_end|>"
95
+ self.merged_str_list.append(message_line)
96
+ self.merged_str_list.append("<|im_start|>assistant")
97
+ self.merged_str = "\n".join(self.merged_str_list)
98
+ # https://huggingface.co/openchat/openchat-3.5-0106
99
+ elif self.model in ["zephyr-7b-beta"]:
100
+ self.messages = self.concat_messages_by_role(messages)
101
+ self.cached_str = ""
102
+ for message in self.messages:
103
+ role = message["role"]
104
+ content = message["content"]
105
+ if role in self.inst_roles:
106
+ self.cached_str = f"[/USER] {content} [/ASSIST]"
107
+ elif role in self.answer_roles:
108
+ self.merged_str += f"<s> {self.cached_str} {content} </s>\n"
109
+ self.cached_str = ""
110
+ else:
111
+ self.cached_str = f"[/USER] {content} [/ASSIST]"
112
+ if self.cached_str:
113
+ self.merged_str += f"{self.cached_str}"
114
+ elif self.model in ["starchat2-15b-v0.1"]:
115
+ self.merged_str_list = []
116
+ for message in self.messages:
117
+ role = message["role"]
118
+ content = message["content"]
119
+ if role not in ["system", "user", "assistant"]:
120
+ role = self.default_role
121
+ message_line = f"<|im_start|>{role}\n{content}<|im_end|>"
122
+ self.merged_str_list.append(message_line)
123
+ self.merged_str_list.append("<|im_start|>assistant")
124
+ self.merged_str = "\n".join(self.merged_str_list)
125
+ else:
126
+ self.merged_str = "\n".join(
127
+ [
128
+ f'`{message["role"]}`:\n{message["content"]}\n'
129
+ for message in self.messages
130
+ ]
131
+ )
132
+
133
+ return self.merged_str
134
+
135
+ def convert_pair_matches_to_messages(self, pair_matches_list):
136
+ messages = []
137
+ if len(pair_matches_list) <= 0:
138
+ messages = [
139
+ {
140
+ "role": "user",
141
+ "content": self.merged_str,
142
+ }
143
+ ]
144
+ else:
145
+ for match in pair_matches_list:
146
+ inst = match.group("inst")
147
+ answer = match.group("answer")
148
+ messages.extend(
149
+ [
150
+ {"role": "user", "content": inst.strip()},
151
+ {"role": "assistant", "content": answer.strip()},
152
+ ]
153
+ )
154
+ return messages
155
+
156
+ def append_last_instruction_to_messages(self, inst_matches_list, pair_matches_list):
157
+ if len(inst_matches_list) > len(pair_matches_list):
158
+ self.messages.extend(
159
+ [
160
+ {
161
+ "role": "user",
162
+ "content": inst_matches_list[-1].group("inst").strip(),
163
+ }
164
+ ]
165
+ )
166
+
167
+ def split(self, merged_str) -> list:
168
+ self.merged_str = merged_str
169
+ self.messages = []
170
+
171
+ if self.model in ["mixtral-8x7b", "mistral-7b"]:
172
+ pair_pattern = (
173
+ r"<s>\s*\[INST\](?P<inst>[\s\S]*?)\[/INST\](?P<answer>[\s\S]*?)</s>"
174
+ )
175
+ pair_matches = re.finditer(pair_pattern, self.merged_str, re.MULTILINE)
176
+ pair_matches_list = list(pair_matches)
177
+
178
+ self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
179
+
180
+ inst_pattern = r"\[INST\](?P<inst>[\s\S]*?)\[/INST\]"
181
+ inst_matches = re.finditer(inst_pattern, self.merged_str, re.MULTILINE)
182
+ inst_matches_list = list(inst_matches)
183
+
184
+ self.append_last_instruction_to_messages(
185
+ inst_matches_list, pair_matches_list
186
+ )
187
+ elif self.model in ["nous-mixtral-8x7b"]:
188
+ # https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO#prompt-format
189
+ # message_pattern = r"<\|im_start\|>(?P<role>system|user|assistant)[\s\n]*(?P<content>[\s\S]*?)<\|im_end\|>"
190
+ message_pattern = r"<\|im_start\|>(?P<role>system|user|assistant)[\s\n]*(?P<content>[\s\S]*?)<\|im_end\|>"
191
+ message_matches = re.finditer(
192
+ message_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
193
+ )
194
+ message_matches_list = list(message_matches)
195
+ logger.note(f"message_matches_list: {message_matches_list}")
196
+ for match in message_matches_list:
197
+ role = match.group("role")
198
+ content = match.group("content")
199
+ self.messages.append({"role": role, "content": content.strip()})
200
+ elif self.model in ["zephyr-7b-beta"]:
201
+ pair_pattern = (
202
+ r"<s>\s*\[INST\](?P<inst>[\s\S]*?)\[/USER\][/ASSIST\](?P<answer>[\s\S]*?)</s>"
203
+ )
204
+ pair_matches = re.finditer(pair_pattern, self.merged_str, re.MULTILINE)
205
+ pair_matches_list = list(pair_matches)
206
+
207
+ self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
208
+
209
+ inst_pattern = r"\[INST\](?P<inst>[\s\S]*?)\[/USER\][/ASSIST\]"
210
+ inst_matches = re.finditer(inst_pattern, self.merged_str, re.MULTILINE)
211
+ inst_matches_list = list(inst_matches)
212
+
213
+ self.append_last_instruction_to_messages(
214
+ inst_matches_list, pair_matches_list
215
+ )
216
+ elif self.model in ["starchat2-15b-v0.1"]:
217
+ message_pattern = r"<\|im_start\|>(?P<role>system|user|assistant)[\s\n]*(?P<content>[\s\S]*?)<\|im_end\|>"
218
+ message_matches = re.finditer(
219
+ message_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
220
+ )
221
+ message_matches_list = list(message_matches)
222
+ logger.note(f"message_matches_list: {message_matches_list}")
223
+ for match in message_matches_list:
224
+ role = match.group("role")
225
+ content = match.group("content")
226
+ self.messages.append({"role": role, "content": content.strip()})
227
+ else:
228
+ self.messages = [
229
+ {
230
+ "role": "user",
231
+ "content": self.merged_str,
232
+ }
233
+ ]
234
+
235
+ return self.messages
236
+
237
+
238
+ if __name__ == "__main__":
239
+ model = "mixtral-8x7b"
240
+ # model = "nous-mixtral-8x7b"
241
+ composer = MessageComposer(model)
242
+ messages = [
243
+ {
244
+ "role": "system",
245
+ "content": "You are a LLM developed by OpenAI.\nYour name is GPT-4.",
246
+ },
247
+ {"role": "user", "content": "Hello, who are you?"},
248
+ {"role": "assistant", "content": "I am a bot."},
249
+ {"role": "user", "content": "What is your name?"},
250
+ # {"role": "assistant", "content": "My name is Bing."},
251
+ # {"role": "user", "content": "Tell me a joke."},
252
+ # {"role": "assistant", "content": "What is a robot's favorite type of music?"},
253
+ # {
254
+ # "role": "user",
255
+ # "content": "How many questions have I asked? Please list them.",
256
+ # },
257
+ ]
258
+ logger.note(f"model: {composer.model}")
259
+ merged_str = composer.merge(messages)
260
+ logger.note("merged_str:")
261
+ logger.mesg(merged_str)
262
+ logger.note("splitted messages:")
263
+ pprint(composer.split(merged_str))
264
+ # logger.note("merged merged_str:")
265
+ # 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,217 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ "zephyr-7b-beta": "HuggingFaceH4/zephyr-7b-beta", # ❌ Too Slow
16
+ "starchat2-15b-v0.1": "HuggingFaceH4/starchat2-15b-v0.1", # ❌ Too Slow
17
+
18
+ # "llama-70b": "meta-llama/Llama-2-70b-chat-hf", # ❌ Require Pro User
19
+ # "codellama-34b": "codellama/CodeLlama-34b-Instruct-hf", # ❌ Low Score
20
+ # "falcon-180b": "tiiuae/falcon-180B-chat", # ❌ Require Pro User
21
+ "default": "mistralai/Mixtral-8x7B-Instruct-v0.1",
22
+ }
23
+ STOP_SEQUENCES_MAP = {
24
+ "mixtral-8x7b": "</s>",
25
+ "mistral-7b": "</s>",
26
+ "nous-mixtral-8x7b": "<|im_end|>",
27
+ "zephyr-7b-beta": "</s>",
28
+ "starchat2-15b-v0.1": "<|im_end|>",
29
+ }
30
+ TOKEN_LIMIT_MAP = {
31
+ "mixtral-8x7b": 32768,
32
+ "mistral-7b": 32768,
33
+ "nous-mixtral-8x7b": 32768,
34
+ "zephyr-7b-beta": 4096,
35
+ "starchat2-15b-v0.1": 8192
36
+
37
+ }
38
+ TOKEN_RESERVED = 100
39
+
40
+ def __init__(self, model: str):
41
+ if model in self.MODEL_MAP.keys():
42
+ self.model = model
43
+ else:
44
+ self.model = "default"
45
+ self.model_fullname = self.MODEL_MAP[self.model]
46
+ self.message_outputer = OpenaiStreamOutputer()
47
+ self.tokenizer = tiktoken_get_encoding("cl100k_base")
48
+
49
+ def parse_line(self, line):
50
+ line = line.decode("utf-8")
51
+ line = re.sub(r"data:\s*", "", line)
52
+ data = json.loads(line)
53
+ try:
54
+ content = data["token"]["text"]
55
+ except:
56
+ logger.err(data)
57
+ return content
58
+
59
+ def count_tokens(self, text):
60
+ tokens = self.tokenizer.encode(text)
61
+ token_count = len(tokens)
62
+ logger.note(f"Prompt Token Count: {token_count}")
63
+ return token_count
64
+
65
+ def chat_response(
66
+ self,
67
+ prompt: str = None,
68
+ temperature: float = 0.5,
69
+ top_p: float = 0.95,
70
+ max_new_tokens: int = None,
71
+ api_key: str = None,
72
+ use_cache: bool = False,
73
+ ):
74
+ # https://huggingface.co/docs/api-inference/detailed_parameters?code=curl
75
+ # 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>'
76
+ self.request_url = (
77
+ f"https://api-inference.huggingface.co/models/{self.model_fullname}"
78
+ )
79
+ self.request_headers = {
80
+ "Content-Type": "application/json",
81
+ }
82
+
83
+ if api_key:
84
+ logger.note(
85
+ f"Using API Key: {api_key[:3]}{(len(api_key)-7)*'*'}{api_key[-4:]}"
86
+ )
87
+ self.request_headers["Authorization"] = f"Bearer {api_key}"
88
+
89
+ if temperature is None or temperature < 0:
90
+ temperature = 0.0
91
+ # temperature must 0 < and < 1 for HF LLM models
92
+ temperature = max(temperature, 0.01)
93
+ temperature = min(temperature, 0.99)
94
+ top_p = max(top_p, 0.01)
95
+ top_p = min(top_p, 0.99)
96
+
97
+ token_limit = int(
98
+ self.TOKEN_LIMIT_MAP[self.model]
99
+ - self.TOKEN_RESERVED
100
+ - self.count_tokens(prompt) * 1.35
101
+ )
102
+ if token_limit <= 0:
103
+ raise ValueError("Prompt exceeded token limit!")
104
+
105
+ if max_new_tokens is None or max_new_tokens <= 0:
106
+ max_new_tokens = token_limit
107
+ else:
108
+ max_new_tokens = min(max_new_tokens, token_limit)
109
+
110
+ # References:
111
+ # huggingface_hub/inference/_client.py:
112
+ # class InferenceClient > def text_generation()
113
+ # huggingface_hub/inference/_text_generation.py:
114
+ # class TextGenerationRequest > param `stream`
115
+ # https://huggingface.co/docs/text-generation-inference/conceptual/streaming#streaming-with-curl
116
+ # https://huggingface.co/docs/api-inference/detailed_parameters#text-generation-task
117
+ self.request_body = {
118
+ "inputs": prompt,
119
+ "parameters": {
120
+ "temperature": temperature,
121
+ "top_p": top_p,
122
+ "max_new_tokens": max_new_tokens,
123
+ "return_full_text": False,
124
+ },
125
+ "options": {
126
+ "use_cache": use_cache,
127
+ },
128
+ "stream": True,
129
+ }
130
+
131
+ if self.model in self.STOP_SEQUENCES_MAP.keys():
132
+ self.stop_sequences = self.STOP_SEQUENCES_MAP[self.model]
133
+ # self.request_body["parameters"]["stop_sequences"] = [
134
+ # self.STOP_SEQUENCES[self.model]
135
+ # ]
136
+
137
+ logger.back(self.request_url)
138
+ enver.set_envs(proxies=True)
139
+ stream_response = requests.post(
140
+ self.request_url,
141
+ headers=self.request_headers,
142
+ json=self.request_body,
143
+ proxies=enver.requests_proxies,
144
+ stream=True,
145
+ )
146
+ status_code = stream_response.status_code
147
+ if status_code == 200:
148
+ logger.success(status_code)
149
+ else:
150
+ logger.err(status_code)
151
+
152
+ return stream_response
153
+
154
+ def chat_return_dict(self, stream_response):
155
+ # https://platform.openai.com/docs/guides/text-generation/chat-completions-response-format
156
+ final_output = self.message_outputer.default_data.copy()
157
+ final_output["choices"] = [
158
+ {
159
+ "index": 0,
160
+ "finish_reason": "stop",
161
+ "message": {
162
+ "role": "assistant",
163
+ "content": "",
164
+ },
165
+ }
166
+ ]
167
+ logger.back(final_output)
168
+
169
+ final_content = ""
170
+ for line in stream_response.iter_lines():
171
+ if not line:
172
+ continue
173
+
174
+ content = self.parse_line(line)
175
+
176
+ if content.strip() == self.stop_sequences:
177
+ logger.success("\n[Finished]")
178
+ break
179
+ else:
180
+ logger.back(content, end="")
181
+ final_content += content
182
+
183
+ if self.model in self.STOP_SEQUENCES_MAP.keys():
184
+ final_content = final_content.replace(self.stop_sequences, "")
185
+
186
+ final_content = final_content.strip()
187
+ final_output["choices"][0]["message"]["content"] = final_content
188
+ return final_output
189
+
190
+ def chat_return_generator(self, stream_response):
191
+ is_finished = False
192
+ line_count = 0
193
+ for line in stream_response.iter_lines():
194
+ if line:
195
+ line_count += 1
196
+ else:
197
+ continue
198
+
199
+ content = self.parse_line(line)
200
+
201
+ if content.strip() == self.stop_sequences:
202
+ content_type = "Finished"
203
+ logger.success("\n[Finished]")
204
+ is_finished = True
205
+ else:
206
+ content_type = "Completions"
207
+ if line_count == 1:
208
+ content = content.lstrip()
209
+ logger.back(content, end="")
210
+
211
+ output = self.message_outputer.output(
212
+ content=content, content_type=content_type
213
+ )
214
+ yield output
215
+
216
+ if not is_finished:
217
+ yield self.message_outputer.output(content="", content_type="Finished")
requirements.txt ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
13
+ passlib
14
+ python-multipart
start_file.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ Restart triggered
test.txt ADDED
File without changes
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)