HF-LLM-API

Dockerfile

@@ -0,0 +1,7 @@
+FROM python:3.11-slim
+WORKDIR $HOME/app
+COPY . .
+RUN pip install -r requirements.txt
+VOLUME /data
+EXPOSE 23333
+CMD ["python", "-m", "apis.chat_api"]

README.md

@@ -1,2 +1,164 @@
-# HF-LLM-API
-Huggingface LLM Inference API in OpenAI message format
+---
+title: HF LLM API
+emoji: ☯️
+colorFrom: gray
+colorTo: gray
+sdk: docker
+app_port: 23333
+---
+
+## HF-LLM-API
+Huggingface LLM Inference API in OpenAI message format.
+
+Project link: https://github.com/Niansuh/HF-LLM-API
+
+## Features
+
+- Available Models (2024/01/22): [#5](https://github.com/Niansuh/HF-LLM-API/issues/2)
+  - `mistral-7b`, `mixtral-8x7b`, `nous-mixtral-8x7b`
+  - Adaptive prompt templates for different models
+- Support OpenAI API format
+  - Enable api endpoint via official `openai-python` package
+- Support both stream and no-stream response
+- Support API Key via both HTTP auth header and env varible [#4](https://github.com/Niansuh/HF-LLM-API/issues/1)
+- Docker deployment
+
+## Run API service
+
+### Run in Command Line
+
+**Install dependencies:**
+
+```bash
+# pipreqs . --force --mode no-pin
+pip install -r requirements.txt
+```
+
+**Run API:**
+
+```bash
+python -m apis.chat_api
+```
+
+## Run via Docker
+
+**Docker build:**
+
+```bash
+sudo docker build -t hf-llm-api:1.0 . --build-arg http_proxy=$http_proxy --build-arg https_proxy=$https_proxy
+```
+
+**Docker run:**
+
+```bash
+# no proxy
+sudo docker run -p 23333:23333 hf-llm-api:1.0
+
+# with proxy
+sudo docker run -p 23333:23333 --env http_proxy="http://<server>:<port>" hf-llm-api:1.0
+```
+
+## API Usage
+
+### Using `openai-python`
+
+See: [`examples/chat_with_openai.py`](https://github.com/Niansuh/HF-LLM-API/blob/main/examples/chat_with_openai.py)
+
+```py
+from openai import OpenAI
+
+# If runnning this service with proxy, you might need to unset `http(s)_proxy`.
+base_url = "http://127.0.0.1:23333"
+# Your own HF_TOKEN
+api_key = "hf_xxxxxxxxxxxxxxxx"
+# use below as non-auth user
+# api_key = "sk-xxx"
+
+client = OpenAI(base_url=base_url, api_key=api_key)
+response = client.chat.completions.create(
+    model="mixtral-8x7b",
+    messages=[
+        {
+            "role": "user",
+            "content": "what is your model",
+        }
+    ],
+    stream=True,
+)
+
+for chunk in response:
+    if chunk.choices[0].delta.content is not None:
+        print(chunk.choices[0].delta.content, end="", flush=True)
+    elif chunk.choices[0].finish_reason == "stop":
+        print()
+    else:
+        pass
+```
+
+### Using post requests
+
+See: [`examples/chat_with_post.py`](https://github.com/Niansuh/HF-LLM-API/blob/main/examples/chat_with_post.py)
+
+
+```py
+import ast
+import httpx
+import json
+import re
+
+# If runnning this service with proxy, you might need to unset `http(s)_proxy`.
+chat_api = "http://127.0.0.1:23333"
+# Your own HF_TOKEN
+api_key = "hf_xxxxxxxxxxxxxxxx"
+# use below as non-auth user
+# api_key = "sk-xxx"
+
+requests_headers = {}
+requests_payload = {
+    "model": "mixtral-8x7b",
+    "messages": [
+        {
+            "role": "user",
+            "content": "what is your model",
+        }
+    ],
+    "stream": True,
+}
+
+with httpx.stream(
+    "POST",
+    chat_api + "/chat/completions",
+    headers=requests_headers,
+    json=requests_payload,
+    timeout=httpx.Timeout(connect=20, read=60, write=20, pool=None),
+) as response:
+    # https://docs.aiohttp.org/en/stable/streams.html
+    # https://github.com/openai/openai-cookbook/blob/main/examples/How_to_stream_completions.ipynb
+    response_content = ""
+    for line in response.iter_lines():
+        remove_patterns = [r"^\s*data:\s*", r"^\s*\[DONE\]\s*"]
+        for pattern in remove_patterns:
+            line = re.sub(pattern, "", line).strip()
+
+        if line:
+            try:
+                line_data = json.loads(line)
+            except Exception as e:
+                try:
+                    line_data = ast.literal_eval(line)
+                except:
+                    print(f"Error: {line}")
+                    raise e
+            # print(f"line: {line_data}")
+            delta_data = line_data["choices"][0]["delta"]
+            finish_reason = line_data["choices"][0]["finish_reason"]
+            if "role" in delta_data:
+                role = delta_data["role"]
+            if "content" in delta_data:
+                delta_content = delta_data["content"]
+                response_content += delta_content
+                print(delta_content, end="", flush=True)
+            if finish_reason == "stop":
+                print()
+
+```

apis/chat_api.py

@@ -0,0 +1,214 @@
+import argparse
+import markdown2
+import os
+import sys
+import uvicorn
+
+from pathlib import Path
+from fastapi import FastAPI, Depends
+from fastapi.responses import HTMLResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import BaseModel, Field
+from typing import Union
+from sse_starlette.sse import EventSourceResponse, ServerSentEvent
+from utils.logger import logger
+from networks.message_streamer import MessageStreamer
+from messagers.message_composer import MessageComposer
+from mocks.stream_chat_mocker import stream_chat_mock
+
+
+class ChatAPIApp:
+    def __init__(self):
+        self.app = FastAPI(
+            docs_url="/",
+            title="HuggingFace LLM API",
+            swagger_ui_parameters={"defaultModelsExpandDepth": -1},
+            version="1.0",
+        )
+        self.setup_routes()
+
+    def get_available_models(self):
+        # https://platform.openai.com/docs/api-reference/models/list
+        # ANCHOR[id=available-models]: Available models
+        self.available_models = {
+            "object": "list",
+            "data": [
+                {
+                    "id": "mixtral-8x7b",
+                    "description": "[mistralai/Mixtral-8x7B-Instruct-v0.1]: https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1",
+                    "object": "model",
+                    "created": 1700000000,
+                    "owned_by": "mistralai",
+                },
+                {
+                    "id": "mistral-7b",
+                    "description": "[mistralai/Mistral-7B-Instruct-v0.2]: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2",
+                    "object": "model",
+                    "created": 1700000000,
+                    "owned_by": "mistralai",
+                },
+                {
+                    "id": "nous-mixtral-8x7b",
+                    "description": "[NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO]: https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
+                    "object": "model",
+                    "created": 1700000000,
+                    "owned_by": "NousResearch",
+                },
+            ],
+        }
+        return self.available_models
+
+    def extract_api_key(
+        credentials: HTTPAuthorizationCredentials = Depends(
+            HTTPBearer(auto_error=False)
+        ),
+    ):
+        api_key = None
+        if credentials:
+            api_key = credentials.credentials
+        else:
+            api_key = os.getenv("HF_TOKEN")
+
+        if api_key:
+            if api_key.startswith("hf_"):
+                return api_key
+            else:
+                logger.warn(f"Invalid HF Token!")
+        else:
+            logger.warn("Not provide HF Token!")
+        return None
+
+    class ChatCompletionsPostItem(BaseModel):
+        model: str = Field(
+            default="mixtral-8x7b",
+            description="(str) `mixtral-8x7b`",
+        )
+        messages: list = Field(
+            default=[{"role": "user", "content": "Hello, who are you?"}],
+            description="(list) Messages",
+        )
+        temperature: Union[float, None] = Field(
+            default=0.5,
+            description="(float) Temperature",
+        )
+        top_p: Union[float, None] = Field(
+            default=0.95,
+            description="(float) top p",
+        )
+        max_tokens: Union[int, None] = Field(
+            default=-1,
+            description="(int) Max tokens",
+        )
+        use_cache: bool = Field(
+            default=False,
+            description="(bool) Use cache",
+        )
+        stream: bool = Field(
+            default=True,
+            description="(bool) Stream",
+        )
+
+    def chat_completions(
+        self, item: ChatCompletionsPostItem, api_key: str = Depends(extract_api_key)
+    ):
+        streamer = MessageStreamer(model=item.model)
+        composer = MessageComposer(model=item.model)
+        composer.merge(messages=item.messages)
+        # streamer.chat = stream_chat_mock
+
+        stream_response = streamer.chat_response(
+            prompt=composer.merged_str,
+            temperature=item.temperature,
+            top_p=item.top_p,
+            max_new_tokens=item.max_tokens,
+            api_key=api_key,
+            use_cache=item.use_cache,
+        )
+        if item.stream:
+            event_source_response = EventSourceResponse(
+                streamer.chat_return_generator(stream_response),
+                media_type="text/event-stream",
+                ping=2000,
+                ping_message_factory=lambda: ServerSentEvent(**{"comment": ""}),
+            )
+            return event_source_response
+        else:
+            data_response = streamer.chat_return_dict(stream_response)
+            return data_response
+
+    def get_readme(self):
+        readme_path = Path(__file__).parents[1] / "README.md"
+        with open(readme_path, "r", encoding="utf-8") as rf:
+            readme_str = rf.read()
+        readme_html = markdown2.markdown(
+            readme_str, extras=["table", "fenced-code-blocks", "highlightjs-lang"]
+        )
+        return readme_html
+
+    def setup_routes(self):
+        for prefix in ["", "/v1", "/api", "/api/v1"]:
+            if prefix in ["/api/v1"]:
+                include_in_schema = True
+            else:
+                include_in_schema = False
+
+            self.app.get(
+                prefix + "/models",
+                summary="Get available models",
+                include_in_schema=include_in_schema,
+            )(self.get_available_models)
+
+            self.app.post(
+                prefix + "/chat/completions",
+                summary="Chat completions in conversation session",
+                include_in_schema=include_in_schema,
+            )(self.chat_completions)
+        self.app.get(
+            "/readme",
+            summary="README of HF LLM API",
+            response_class=HTMLResponse,
+            include_in_schema=False,
+        )(self.get_readme)
+
+
+class ArgParser(argparse.ArgumentParser):
+    def __init__(self, *args, **kwargs):
+        super(ArgParser, self).__init__(*args, **kwargs)
+
+        self.add_argument(
+            "-s",
+            "--server",
+            type=str,
+            default="0.0.0.0",
+            help="Server IP for HF LLM Chat API",
+        )
+        self.add_argument(
+            "-p",
+            "--port",
+            type=int,
+            default=23333,
+            help="Server Port for HF LLM Chat API",
+        )
+
+        self.add_argument(
+            "-d",
+            "--dev",
+            default=False,
+            action="store_true",
+            help="Run in dev mode",
+        )
+
+        self.args = self.parse_args(sys.argv[1:])
+
+
+app = ChatAPIApp().app
+
+if __name__ == "__main__":
+    args = ArgParser().args
+    if args.dev:
+        uvicorn.run("__main__:app", host=args.server, port=args.port, reload=True)
+    else:
+        uvicorn.run("__main__:app", host=args.server, port=args.port, reload=False)
+
+    # python -m apis.chat_api      # [Docker] on product mode
+    # python -m apis.chat_api -d   # [Dev]    on develop mode

examples/chat_with_openai.py

@@ -0,0 +1,25 @@
+from openai import OpenAI
+
+# If runnning this service with proxy, you might need to unset `http(s)_proxy`.
+base_url = "http://127.0.0.1:23333"
+api_key = "sk-xxxxx"
+
+client = OpenAI(base_url=base_url, api_key=api_key)
+response = client.chat.completions.create(
+    model="mixtral-8x7b",
+    messages=[
+        {
+            "role": "user",
+            "content": "what is your model",
+        }
+    ],
+    stream=True,
+)
+
+for chunk in response:
+    if chunk.choices[0].delta.content is not None:
+        print(chunk.choices[0].delta.content, end="", flush=True)
+    elif chunk.choices[0].finish_reason == "stop":
+        print()
+    else:
+        pass

examples/chat_with_post.py

@@ -0,0 +1,55 @@
+import ast
+import httpx
+import json
+import re
+
+# If runnning this service with proxy, you might need to unset `http(s)_proxy`.
+chat_api = "http://127.0.0.1:23333"
+api_key = "sk-xxxxx"
+requests_headers = {}
+requests_payload = {
+    "model": "mixtral-8x7b",
+    "messages": [
+        {
+            "role": "user",
+            "content": "what is your model",
+        }
+    ],
+    "stream": True,
+}
+
+with httpx.stream(
+    "POST",
+    chat_api + "/chat/completions",
+    headers=requests_headers,
+    json=requests_payload,
+    timeout=httpx.Timeout(connect=20, read=60, write=20, pool=None),
+) as response:
+    # https://docs.aiohttp.org/en/stable/streams.html
+    # https://github.com/openai/openai-cookbook/blob/main/examples/How_to_stream_completions.ipynb
+    response_content = ""
+    for line in response.iter_lines():
+        remove_patterns = [r"^\s*data:\s*", r"^\s*\[DONE\]\s*"]
+        for pattern in remove_patterns:
+            line = re.sub(pattern, "", line).strip()
+
+        if line:
+            try:
+                line_data = json.loads(line)
+            except Exception as e:
+                try:
+                    line_data = ast.literal_eval(line)
+                except:
+                    print(f"Error: {line}")
+                    raise e
+            # print(f"line: {line_data}")
+            delta_data = line_data["choices"][0]["delta"]
+            finish_reason = line_data["choices"][0]["finish_reason"]
+            if "role" in delta_data:
+                role = delta_data["role"]
+            if "content" in delta_data:
+                delta_content = delta_data["content"]
+                response_content += delta_content
+                print(delta_content, end="", flush=True)
+            if finish_reason == "stop":
+                print()

messagers/message_composer.py

@@ -0,0 +1,247 @@
+import re
+from pprint import pprint
+from utils.logger import logger
+
+
+class MessageComposer:
+    # LINK - apis/chat_api.py#available-models
+    AVALAIBLE_MODELS = [
+        "mixtral-8x7b",
+        "mistral-7b",
+        "openchat-3.5",
+        "nous-mixtral-8x7b",
+    ]
+
+    def __init__(self, model: str = None):
+        if model in self.AVALAIBLE_MODELS:
+            self.model = model
+        else:
+            self.model = "mixtral-8x7b"
+        self.system_roles = ["system"]
+        self.inst_roles = ["user", "system", "inst"]
+        self.answer_roles = ["assistant", "bot", "answer"]
+        self.default_role = "user"
+
+    def concat_messages_by_role(self, messages):
+        def is_same_role(role1, role2):
+            if (
+                (role1 == role2)
+                or (role1 in self.inst_roles and role2 in self.inst_roles)
+                or (role1 in self.answer_roles and role2 in self.answer_roles)
+            ):
+                return True
+            else:
+                return False
+
+        concat_messages = []
+        for message in messages:
+            role = message["role"]
+            content = message["content"]
+            if concat_messages and is_same_role(role, concat_messages[-1]["role"]):
+                concat_messages[-1]["content"] += "\n" + content
+            else:
+                if role in self.inst_roles:
+                    message["role"] = "inst"
+                elif role in self.answer_roles:
+                    message["role"] = "answer"
+                else:
+                    message["role"] = "inst"
+                concat_messages.append(message)
+        return concat_messages
+
+    def merge(self, messages) -> str:
+        # Mistral and Mixtral:
+        #   <s> [INST] Instruction [/INST] Model answer </s> [INST] Follow-up instruction [/INST]
+
+        # OpenChat:
+        #   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:
+
+        # Nous Mixtral:
+        #   <|im_start|>system
+        #   You are "Hermes 2".<|im_end|>
+        #   <|im_start|>user
+        #   Hello, who are you?<|im_end|>
+        #   <|im_start|>assistant
+
+        self.messages = messages
+        self.merged_str = ""
+
+        # https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1#instruction-format
+        if self.model in ["mixtral-8x7b", "mistral-7b"]:
+            self.messages = self.concat_messages_by_role(messages)
+            self.cached_str = ""
+            for message in self.messages:
+                role = message["role"]
+                content = message["content"]
+                if role in self.inst_roles:
+                    self.cached_str = f"[INST] {content} [/INST]"
+                elif role in self.answer_roles:
+                    self.merged_str += f"<s> {self.cached_str} {content} </s>\n"
+                    self.cached_str = ""
+                else:
+                    self.cached_str = f"[INST] {content} [/INST]"
+            if self.cached_str:
+                self.merged_str += f"{self.cached_str}"
+        # https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO#prompt-format
+        elif self.model in ["nous-mixtral-8x7b"]:
+            self.merged_str_list = []
+            for message in self.messages:
+                role = message["role"]
+                content = message["content"]
+                if role not in ["system", "user", "assistant"]:
+                    role = self.default_role
+                message_line = f"<|im_start|>{role}\n{content}<|im_end|>"
+                self.merged_str_list.append(message_line)
+            self.merged_str_list.append("<|im_start|>assistant")
+            self.merged_str = "\n".join(self.merged_str_list)
+        # https://huggingface.co/openchat/openchat-3.5-0106
+        elif self.model in ["openchat-3.5"]:
+            self.messages = self.concat_messages_by_role(messages)
+            self.merged_str_list = []
+            self.end_of_turn = "<|end_of_turn|>"
+            for message in self.messages:
+                role = message["role"]
+                content = message["content"]
+                if role in self.inst_roles:
+                    self.merged_str_list.append(
+                        f"GPT4 Correct User:\n{content}{self.end_of_turn}"
+                    )
+                elif role in self.answer_roles:
+                    self.merged_str_list.append(
+                        f"GPT4 Correct Assistant:\n{content}{self.end_of_turn}"
+                    )
+                else:
+                    self.merged_str_list.append(
+                        f"GPT4 Correct User: {content}{self.end_of_turn}"
+                    )
+            self.merged_str_list.append(f"GPT4 Correct Assistant:\n")
+            self.merged_str = "\n".join(self.merged_str_list)
+        else:
+            self.merged_str = "\n".join(
+                [
+                    f'`{message["role"]}`:\n{message["content"]}\n'
+                    for message in self.messages
+                ]
+            )
+
+        return self.merged_str
+
+    def convert_pair_matches_to_messages(self, pair_matches_list):
+        messages = []
+        if len(pair_matches_list) <= 0:
+            messages = [
+                {
+                    "role": "user",
+                    "content": self.merged_str,
+                }
+            ]
+        else:
+            for match in pair_matches_list:
+                inst = match.group("inst")
+                answer = match.group("answer")
+                messages.extend(
+                    [
+                        {"role": "user", "content": inst.strip()},
+                        {"role": "assistant", "content": answer.strip()},
+                    ]
+                )
+        return messages
+
+    def append_last_instruction_to_messages(self, inst_matches_list, pair_matches_list):
+        if len(inst_matches_list) > len(pair_matches_list):
+            self.messages.extend(
+                [
+                    {
+                        "role": "user",
+                        "content": inst_matches_list[-1].group("inst").strip(),
+                    }
+                ]
+            )
+
+    def split(self, merged_str) -> list:
+        self.merged_str = merged_str
+        self.messages = []
+
+        if self.model in ["mixtral-8x7b", "mistral-7b"]:
+            pair_pattern = (
+                r"<s>\s*\[INST\](?P<inst>[\s\S]*?)\[/INST\](?P<answer>[\s\S]*?)</s>"
+            )
+            pair_matches = re.finditer(pair_pattern, self.merged_str, re.MULTILINE)
+            pair_matches_list = list(pair_matches)
+
+            self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
+
+            inst_pattern = r"\[INST\](?P<inst>[\s\S]*?)\[/INST\]"
+            inst_matches = re.finditer(inst_pattern, self.merged_str, re.MULTILINE)
+            inst_matches_list = list(inst_matches)
+
+            self.append_last_instruction_to_messages(
+                inst_matches_list, pair_matches_list
+            )
+        elif self.model in ["nous-mixtral-8x7b"]:
+            # https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO#prompt-format
+            # message_pattern = r"<\|im_start\|>(?P<role>system|user|assistant)[\s\n]*(?P<content>[\s\S]*?)<\|im_end\|>"
+            message_pattern = r"<\|im_start\|>(?P<role>system|user|assistant)[\s\n]*(?P<content>[\s\S]*?)<\|im_end\|>"
+            message_matches = re.finditer(
+                message_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
+            )
+            message_matches_list = list(message_matches)
+            logger.note(f"message_matches_list: {message_matches_list}")
+            for match in message_matches_list:
+                role = match.group("role")
+                content = match.group("content")
+                self.messages.append({"role": role, "content": content.strip()})
+        elif self.model in ["openchat-3.5"]:
+            pair_pattern = r"GPT4 Correct User:(?P<inst>[\s\S]*?)<\|end_of_turn\|>\s*GPT4 Correct Assistant:(?P<answer>[\s\S]*?)<\|end_of_turn\|>"
+            pair_matches = re.finditer(
+                pair_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
+            )
+            pair_matches_list = list(pair_matches)
+            self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
+            inst_pattern = r"GPT4 Correct User:(?P<inst>[\s\S]*?)<\|end_of_turn\|>"
+            inst_matches = re.finditer(
+                inst_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
+            )
+            inst_matches_list = list(inst_matches)
+            self.append_last_instruction_to_messages(
+                inst_matches_list, pair_matches_list
+            )
+        else:
+            self.messages = [
+                {
+                    "role": "user",
+                    "content": self.merged_str,
+                }
+            ]
+
+        return self.messages
+
+
+if __name__ == "__main__":
+    model = "mixtral-8x7b"
+    # model = "nous-mixtral-8x7b"
+    composer = MessageComposer(model)
+    messages = [
+        {
+            "role": "system",
+            "content": "You are a LLM developed by OpenAI.\nYour name is GPT-4.",
+        },
+        {"role": "user", "content": "Hello, who are you?"},
+        {"role": "assistant", "content": "I am a bot."},
+        {"role": "user", "content": "What is your name?"},
+        # {"role": "assistant", "content": "My name is Bing."},
+        # {"role": "user", "content": "Tell me a joke."},
+        # {"role": "assistant", "content": "What is a robot's favorite type of music?"},
+        # {
+        #     "role": "user",
+        #     "content": "How many questions have I asked? Please list them.",
+        # },
+    ]
+    logger.note(f"model: {composer.model}")
+    merged_str = composer.merge(messages)
+    logger.note("merged_str:")
+    logger.mesg(merged_str)
+    logger.note("splitted messages:")
+    pprint(composer.split(merged_str))
+    # logger.note("merged merged_str:")
+    # logger.mesg(composer.merge(composer.split(merged_str)))

messagers/message_outputer.py

@@ -0,0 +1,65 @@
+import json
+
+
+class OpenaiStreamOutputer:
+    """
+    Create chat completion - OpenAI API Documentation
+    * https://platform.openai.com/docs/api-reference/chat/create
+    """
+
+    def __init__(self):
+        self.default_data = {
+            "created": 1700000000,
+            "id": "chatcmpl-hugginface",
+            "object": "chat.completion.chunk",
+            # "content_type": "Completions",
+            "model": "hugginface",
+            "choices": [],
+        }
+
+    def data_to_string(self, data={}, content_type=""):
+        data_str = f"{json.dumps(data)}"
+        return data_str
+
+    def output(self, content=None, content_type="Completions") -> str:
+        data = self.default_data.copy()
+        if content_type == "Role":
+            data["choices"] = [
+                {
+                    "index": 0,
+                    "delta": {"role": "assistant"},
+                    "finish_reason": None,
+                }
+            ]
+        elif content_type in [
+            "Completions",
+            "InternalSearchQuery",
+            "InternalSearchResult",
+            "SuggestedResponses",
+        ]:
+            if content_type in ["InternalSearchQuery", "InternalSearchResult"]:
+                content += "\n"
+            data["choices"] = [
+                {
+                    "index": 0,
+                    "delta": {"content": content},
+                    "finish_reason": None,
+                }
+            ]
+        elif content_type == "Finished":
+            data["choices"] = [
+                {
+                    "index": 0,
+                    "delta": {},
+                    "finish_reason": "stop",
+                }
+            ]
+        else:
+            data["choices"] = [
+                {
+                    "index": 0,
+                    "delta": {},
+                    "finish_reason": None,
+                }
+            ]
+        return self.data_to_string(data, content_type)

mocks/stream_chat_mocker.py

@@ -0,0 +1,13 @@
+import time
+from utils.logger import logger
+
+
+def stream_chat_mock(*args, **kwargs):
+    logger.note(msg=str(args) + str(kwargs))
+    for i in range(10):
+        content = f"W{i+1} "
+        time.sleep(0.1)
+        logger.mesg(content, end="")
+        yield content
+    logger.mesg("")
+    yield ""

networks/message_streamer.py

@@ -0,0 +1,212 @@
+import json
+import re
+import requests
+from tiktoken import get_encoding as tiktoken_get_encoding
+from messagers.message_outputer import OpenaiStreamOutputer
+from utils.logger import logger
+from utils.enver import enver
+
+
+class MessageStreamer:
+    MODEL_MAP = {
+        "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1",  # 72.62, fast [Recommended]
+        "mistral-7b": "mistralai/Mistral-7B-Instruct-v0.2",  # 65.71, fast
+        "nous-mixtral-8x7b": "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
+        # "openchat-3.5": "openchat/openchat-3.5-1210",  # 68.89, fast
+        # "zephyr-7b-beta": "HuggingFaceH4/zephyr-7b-beta",  # ❌ Too Slow
+        # "llama-70b": "meta-llama/Llama-2-70b-chat-hf",  # ❌ Require Pro User
+        # "codellama-34b": "codellama/CodeLlama-34b-Instruct-hf",  # ❌ Low Score
+        # "falcon-180b": "tiiuae/falcon-180B-chat",  # ❌ Require Pro User
+        "default": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    }
+    STOP_SEQUENCES_MAP = {
+        "mixtral-8x7b": "</s>",
+        "mistral-7b": "</s>",
+        "nous-mixtral-8x7b": "<|im_end|>",
+        "openchat-3.5": "<|end_of_turn|>",
+    }
+    TOKEN_LIMIT_MAP = {
+        "mixtral-8x7b": 32768,
+        "mistral-7b": 32768,
+        "nous-mixtral-8x7b": 32768,
+        "openchat-3.5": 8192,
+    }
+    TOKEN_RESERVED = 100
+
+    def __init__(self, model: str):
+        if model in self.MODEL_MAP.keys():
+            self.model = model
+        else:
+            self.model = "default"
+        self.model_fullname = self.MODEL_MAP[self.model]
+        self.message_outputer = OpenaiStreamOutputer()
+        self.tokenizer = tiktoken_get_encoding("cl100k_base")
+
+    def parse_line(self, line):
+        line = line.decode("utf-8")
+        line = re.sub(r"data:\s*", "", line)
+        data = json.loads(line)
+        try:
+            content = data["token"]["text"]
+        except:
+            logger.err(data)
+        return content
+
+    def count_tokens(self, text):
+        tokens = self.tokenizer.encode(text)
+        token_count = len(tokens)
+        logger.note(f"Prompt Token Count: {token_count}")
+        return token_count
+
+    def chat_response(
+        self,
+        prompt: str = None,
+        temperature: float = 0.5,
+        top_p: float = 0.95,
+        max_new_tokens: int = None,
+        api_key: str = None,
+        use_cache: bool = False,
+    ):
+        # https://huggingface.co/docs/api-inference/detailed_parameters?code=curl
+        # 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>'
+        self.request_url = (
+            f"https://api-inference.huggingface.co/models/{self.model_fullname}"
+        )
+        self.request_headers = {
+            "Content-Type": "application/json",
+        }
+
+        if api_key:
+            logger.note(
+                f"Using API Key: {api_key[:3]}{(len(api_key)-7)*'*'}{api_key[-4:]}"
+            )
+            self.request_headers["Authorization"] = f"Bearer {api_key}"
+
+        if temperature is None or temperature < 0:
+            temperature = 0.0
+        # temperature must  0 < and < 1 for HF LLM models
+        temperature = max(temperature, 0.01)
+        temperature = min(temperature, 0.99)
+        top_p = max(top_p, 0.01)
+        top_p = min(top_p, 0.99)
+
+        token_limit = int(
+            self.TOKEN_LIMIT_MAP[self.model]
+            - self.TOKEN_RESERVED
+            - self.count_tokens(prompt) * 1.35
+        )
+        if token_limit <= 0:
+            raise ValueError("Prompt exceeded token limit!")
+
+        if max_new_tokens is None or max_new_tokens <= 0:
+            max_new_tokens = token_limit
+        else:
+            max_new_tokens = min(max_new_tokens, token_limit)
+
+        # References:
+        #   huggingface_hub/inference/_client.py:
+        #     class InferenceClient > def text_generation()
+        #   huggingface_hub/inference/_text_generation.py:
+        #     class TextGenerationRequest > param `stream`
+        # https://huggingface.co/docs/text-generation-inference/conceptual/streaming#streaming-with-curl
+        # https://huggingface.co/docs/api-inference/detailed_parameters#text-generation-task
+        self.request_body = {
+            "inputs": prompt,
+            "parameters": {
+                "temperature": temperature,
+                "top_p": top_p,
+                "max_new_tokens": max_new_tokens,
+                "return_full_text": False,
+            },
+            "options": {
+                "use_cache": use_cache,
+            },
+            "stream": True,
+        }
+
+        if self.model in self.STOP_SEQUENCES_MAP.keys():
+            self.stop_sequences = self.STOP_SEQUENCES_MAP[self.model]
+        #     self.request_body["parameters"]["stop_sequences"] = [
+        #         self.STOP_SEQUENCES[self.model]
+        #     ]
+
+        logger.back(self.request_url)
+        enver.set_envs(proxies=True)
+        stream_response = requests.post(
+            self.request_url,
+            headers=self.request_headers,
+            json=self.request_body,
+            proxies=enver.requests_proxies,
+            stream=True,
+        )
+        status_code = stream_response.status_code
+        if status_code == 200:
+            logger.success(status_code)
+        else:
+            logger.err(status_code)
+
+        return stream_response
+
+    def chat_return_dict(self, stream_response):
+        # https://platform.openai.com/docs/guides/text-generation/chat-completions-response-format
+        final_output = self.message_outputer.default_data.copy()
+        final_output["choices"] = [
+            {
+                "index": 0,
+                "finish_reason": "stop",
+                "message": {
+                    "role": "assistant",
+                    "content": "",
+                },
+            }
+        ]
+        logger.back(final_output)
+
+        final_content = ""
+        for line in stream_response.iter_lines():
+            if not line:
+                continue
+            content = self.parse_line(line)
+
+            if content.strip() == self.stop_sequences:
+                logger.success("\n[Finished]")
+                break
+            else:
+                logger.back(content, end="")
+                final_content += content
+
+        if self.model in self.STOP_SEQUENCES_MAP.keys():
+            final_content = final_content.replace(self.stop_sequences, "")
+
+        final_content = final_content.strip()
+        final_output["choices"][0]["message"]["content"] = final_content
+        return final_output
+
+    def chat_return_generator(self, stream_response):
+        is_finished = False
+        line_count = 0
+        for line in stream_response.iter_lines():
+            if line:
+                line_count += 1
+            else:
+                continue
+
+            content = self.parse_line(line)
+
+            if content.strip() == self.stop_sequences:
+                content_type = "Finished"
+                logger.success("\n[Finished]")
+                is_finished = True
+            else:
+                content_type = "Completions"
+                if line_count == 1:
+                    content = content.lstrip()
+                logger.back(content, end="")
+
+            output = self.message_outputer.output(
+                content=content, content_type=content_type
+            )
+            yield output
+
+        if not is_finished:
+            yield self.message_outputer.output(content="", content_type="Finished")

requirements.txt

@@ -0,0 +1,12 @@
+aiohttp
+fastapi
+httpx
+markdown2[all]
+openai
+pydantic
+requests
+sse_starlette
+termcolor
+tiktoken
+uvicorn
+websockets

utils/__init__.py

@@ -0,0 +1,69 @@
+import json
+import requests
+import os
+
+from pathlib import Path
+
+
+class OSEnver:
+    def __init__(self):
+        self.envs_stack = []
+        self.envs = os.environ.copy()
+
+    def store_envs(self):
+        self.envs_stack.append(self.envs)
+
+    def restore_envs(self):
+        self.envs = self.envs_stack.pop()
+        if self.global_scope:
+            os.environ = self.envs
+
+    def set_envs(self, secrets=True, proxies=None, store_envs=True):
+        # caller_info = inspect.stack()[1]
+        # logger.back(f"OS Envs is set by: {caller_info.filename}")
+
+        if store_envs:
+            self.store_envs()
+
+        if secrets:
+            secrets_path = Path(__file__).parents[1] / "secrets.json"
+            if secrets_path.exists():
+                with open(secrets_path, "r") as rf:
+                    secrets = json.load(rf)
+            else:
+                secrets = {}
+
+        if proxies:
+            for proxy_env in ["http_proxy", "https_proxy"]:
+                if isinstance(proxies, str):
+                    self.envs[proxy_env] = proxies
+                elif "http_proxy" in secrets.keys():
+                    self.envs[proxy_env] = secrets["http_proxy"]
+                elif os.getenv("http_proxy"):
+                    self.envs[proxy_env] = os.getenv("http_proxy")
+                else:
+                    continue
+
+        self.proxy = (
+            self.envs.get("all_proxy")
+            or self.envs.get("http_proxy")
+            or self.envs.get("https_proxy")
+            or None
+        )
+        self.requests_proxies = {
+            "http": self.proxy,
+            "https": self.proxy,
+        }
+
+        # https://www.proxynova.com/proxy-server-list/country-us/
+
+        print(f"Using proxy: [{self.proxy}]")
+        # r = requests.get(
+        #     "http://ifconfig.me/ip",
+        #     proxies=self.requests_proxies,
+        #     timeout=10,
+        # )
+        # print(f"[r.status_code] r.text")
+
+
+enver = OSEnver()

utils/enver.py

@@ -0,0 +1,60 @@
+import json
+import os
+
+from pathlib import Path
+from utils.logger import logger
+
+
+class OSEnver:
+    def __init__(self):
+        self.envs_stack = []
+        self.envs = os.environ.copy()
+
+    def store_envs(self):
+        self.envs_stack.append(self.envs)
+
+    def restore_envs(self):
+        self.envs = self.envs_stack.pop()
+
+    def set_envs(self, secrets=True, proxies=None, store_envs=True):
+        # caller_info = inspect.stack()[1]
+        # logger.back(f"OS Envs is set by: {caller_info.filename}")
+
+        if store_envs:
+            self.store_envs()
+
+        if secrets:
+            secrets_path = Path(__file__).parents[1] / "secrets.json"
+            if secrets_path.exists():
+                with open(secrets_path, "r") as rf:
+                    secrets = json.load(rf)
+            else:
+                secrets = {}
+
+        if proxies:
+            for proxy_env in ["http_proxy", "https_proxy"]:
+                if isinstance(proxies, str):
+                    self.envs[proxy_env] = proxies
+                elif "http_proxy" in secrets.keys():
+                    self.envs[proxy_env] = secrets["http_proxy"]
+                elif os.getenv("http_proxy"):
+                    self.envs[proxy_env] = os.getenv("http_proxy")
+                else:
+                    continue
+
+        self.proxy = (
+            self.envs.get("all_proxy")
+            or self.envs.get("http_proxy")
+            or self.envs.get("https_proxy")
+            or None
+        )
+        self.requests_proxies = {
+            "http": self.proxy,
+            "https": self.proxy,
+        }
+
+        if self.proxy:
+            logger.note(f"Using proxy: [{self.proxy}]")
+
+
+enver = OSEnver()

utils/logger.py

@@ -0,0 +1,269 @@
+import datetime
+import functools
+import inspect
+import logging
+import os
+import shutil
+import subprocess
+from termcolor import colored
+
+
+def add_fillers(text, filler="=", fill_side="both"):
+    terminal_width = shutil.get_terminal_size().columns
+    text = text.strip()
+    text_width = len(text)
+    if text_width >= terminal_width:
+        return text
+
+    if fill_side[0].lower() == "b":
+        leading_fill_str = filler * ((terminal_width - text_width) // 2 - 1) + " "
+        trailing_fill_str = " " + filler * (
+            terminal_width - text_width - len(leading_fill_str) - 1
+        )
+    elif fill_side[0].lower() == "l":
+        leading_fill_str = filler * (terminal_width - text_width - 1) + " "
+        trailing_fill_str = ""
+    elif fill_side[0].lower() == "r":
+        leading_fill_str = ""
+        trailing_fill_str = " " + filler * (terminal_width - text_width - 1)
+    else:
+        raise ValueError("Invalid fill_side")
+
+    filled_str = f"{leading_fill_str}{text}{trailing_fill_str}"
+    return filled_str
+
+
+class OSLogger(logging.Logger):
+    LOG_METHODS = {
+        "err": ("error", "red"),
+        "warn": ("warning", "light_red"),
+        "note": ("info", "light_magenta"),
+        "mesg": ("info", "light_cyan"),
+        "file": ("info", "light_blue"),
+        "line": ("info", "white"),
+        "success": ("info", "light_green"),
+        "fail": ("info", "light_red"),
+        "back": ("debug", "light_cyan"),
+    }
+    INDENT_METHODS = [
+        "indent",
+        "set_indent",
+        "reset_indent",
+        "store_indent",
+        "restore_indent",
+        "log_indent",
+    ]
+    LEVEL_METHODS = [
+        "set_level",
+        "store_level",
+        "restore_level",
+        "quiet",
+        "enter_quiet",
+        "exit_quiet",
+    ]
+    LEVEL_NAMES = {
+        "critical": logging.CRITICAL,
+        "error": logging.ERROR,
+        "warning": logging.WARNING,
+        "info": logging.INFO,
+        "debug": logging.DEBUG,
+    }
+
+    def __init__(self, name=None, prefix=False):
+        if not name:
+            frame = inspect.stack()[1]
+            module = inspect.getmodule(frame[0])
+            name = module.__name__
+
+        super().__init__(name)
+        self.setLevel(logging.INFO)
+
+        if prefix:
+            formatter_prefix = "[%(asctime)s] - [%(name)s] - [%(levelname)s]\n"
+        else:
+            formatter_prefix = ""
+
+        self.formatter = logging.Formatter(formatter_prefix + "%(message)s")
+
+        stream_handler = logging.StreamHandler()
+        stream_handler.setLevel(logging.INFO)
+        stream_handler.setFormatter(self.formatter)
+        self.addHandler(stream_handler)
+
+        self.log_indent = 0
+        self.log_indents = []
+
+        self.log_level = "info"
+        self.log_levels = []
+
+    def indent(self, indent=2):
+        self.log_indent += indent
+
+    def set_indent(self, indent=2):
+        self.log_indent = indent
+
+    def reset_indent(self):
+        self.log_indent = 0
+
+    def store_indent(self):
+        self.log_indents.append(self.log_indent)
+
+    def restore_indent(self):
+        self.log_indent = self.log_indents.pop(-1)
+
+    def set_level(self, level):
+        self.log_level = level
+        self.setLevel(self.LEVEL_NAMES[level])
+
+    def store_level(self):
+        self.log_levels.append(self.log_level)
+
+    def restore_level(self):
+        self.log_level = self.log_levels.pop(-1)
+        self.set_level(self.log_level)
+
+    def quiet(self):
+        self.set_level("critical")
+
+    def enter_quiet(self, quiet=False):
+        if quiet:
+            self.store_level()
+            self.quiet()
+
+    def exit_quiet(self, quiet=False):
+        if quiet:
+            self.restore_level()
+
+    def log(
+        self,
+        level,
+        color,
+        msg,
+        indent=0,
+        fill=False,
+        fill_side="both",
+        end="\n",
+        *args,
+        **kwargs,
+    ):
+        if type(msg) == str:
+            msg_str = msg
+        else:
+            msg_str = repr(msg)
+            quotes = ["'", '"']
+            if msg_str[0] in quotes and msg_str[-1] in quotes:
+                msg_str = msg_str[1:-1]
+
+        indent_str = " " * (self.log_indent + indent)
+        indented_msg = "\n".join([indent_str + line for line in msg_str.split("\n")])
+
+        if fill:
+            indented_msg = add_fillers(indented_msg, fill_side=fill_side)
+
+        handler = self.handlers[0]
+        handler.terminator = end
+
+        getattr(self, level)(colored(indented_msg, color), *args, **kwargs)
+
+    def route_log(self, method, msg, *args, **kwargs):
+        level, method = method
+        functools.partial(self.log, level, method, msg)(*args, **kwargs)
+
+    def err(self, msg: str = "", *args, **kwargs):
+        self.route_log(("error", "red"), msg, *args, **kwargs)
+
+    def warn(self, msg: str = "", *args, **kwargs):
+        self.route_log(("warning", "light_red"), msg, *args, **kwargs)
+
+    def note(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "light_magenta"), msg, *args, **kwargs)
+
+    def mesg(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "light_cyan"), msg, *args, **kwargs)
+
+    def file(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "light_blue"), msg, *args, **kwargs)
+
+    def line(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "white"), msg, *args, **kwargs)
+
+    def success(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "light_green"), msg, *args, **kwargs)
+
+    def fail(self, msg: str = "", *args, **kwargs):
+        self.route_log(("info", "light_red"), msg, *args, **kwargs)
+
+    def back(self, msg: str = "", *args, **kwargs):
+        self.route_log(("debug", "light_cyan"), msg, *args, **kwargs)
+
+
+logger = OSLogger()
+
+
+def shell_cmd(cmd, getoutput=False, showcmd=True, env=None):
+    if showcmd:
+        logger.info(colored(f"\n$ [{os.getcwd()}]", "light_blue"))
+        logger.info(colored(f"  $ {cmd}\n", "light_cyan"))
+    if getoutput:
+        output = subprocess.getoutput(cmd, env=env)
+        return output
+    else:
+        subprocess.run(cmd, shell=True, env=env)
+
+
+class Runtimer:
+    def __enter__(self):
+        self.t1, _ = self.start_time()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.t2, _ = self.end_time()
+        self.elapsed_time(self.t2 - self.t1)
+
+    def start_time(self):
+        t1 = datetime.datetime.now()
+        self.logger_time("start", t1)
+        return t1, self.time2str(t1)
+
+    def end_time(self):
+        t2 = datetime.datetime.now()
+        self.logger_time("end", t2)
+        return t2, self.time2str(t2)
+
+    def elapsed_time(self, dt=None):
+        if dt is None:
+            dt = self.t2 - self.t1
+        self.logger_time("elapsed", dt)
+        return dt, self.time2str(dt)
+
+    def logger_time(self, time_type, t):
+        time_types = {
+            "start": "Start",
+            "end": "End",
+            "elapsed": "Elapsed",
+        }
+        time_str = add_fillers(
+            colored(
+                f"{time_types[time_type]} time: [ {self.time2str(t)} ]",
+                "light_magenta",
+            ),
+            fill_side="both",
+        )
+        logger.line(time_str)
+
+    # Convert time to string
+    def time2str(self, t):
+        datetime_str_format = "%Y-%m-%d %H:%M:%S"
+        if isinstance(t, datetime.datetime):
+            return t.strftime(datetime_str_format)
+        elif isinstance(t, datetime.timedelta):
+            hours = t.seconds // 3600
+            hour_str = f"{hours} hr" if hours > 0 else ""
+            minutes = (t.seconds // 60) % 60
+            minute_str = f"{minutes:>2} min" if minutes > 0 else ""
+            seconds = t.seconds % 60
+            second_str = f"{seconds:>2} s"
+            time_str = " ".join([hour_str, minute_str, second_str]).strip()
+            return time_str
+        else:
+            return str(t)