Spaces:
Sleeping
Sleeping
File size: 3,577 Bytes
9f341cc e916990 6aa8b86 9f341cc 6aa8b86 e820e51 9f341cc 6aa8b86 9f341cc e820e51 9f341cc 395ee29 e820e51 9f341cc 3dfcb72 9f341cc b96cef7 9f341cc 047008b 9f341cc 395ee29 9f341cc 047008b 9f341cc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import json
import re
import requests
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
"openchat-3.5": "openchat/openchat_3.5", # 61.24, fast
# "zephyr-7b-alpha": "HuggingFaceH4/zephyr-7b-alpha", # 59.5, fast
# "zephyr-7b-beta": "HuggingFaceH4/zephyr-7b-beta", # 61.95, slow
"default": "mistralai/Mixtral-8x7B-Instruct-v0.1",
}
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()
def parse_line(self, line):
line = line.decode("utf-8")
line = re.sub(r"data:\s*", "", line)
data = json.loads(line)
content = data["token"]["text"]
return content
def chat(
self,
prompt: str = None,
temperature: float = 0.01,
max_new_tokens: int = 8192,
stream: bool = True,
yield_output: 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",
}
# 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
self.request_body = {
"inputs": prompt,
"parameters": {
"temperature": max(temperature, 0.01), # must be positive
"max_new_tokens": max_new_tokens,
"return_full_text": False,
},
"stream": stream,
}
logger.back(self.request_url)
enver.set_envs(proxies=True)
stream = requests.post(
self.request_url,
headers=self.request_headers,
json=self.request_body,
proxies=enver.requests_proxies,
stream=stream,
)
status_code = stream.status_code
if status_code == 200:
logger.success(status_code)
else:
logger.err(status_code)
for line in stream.iter_lines():
if not line:
continue
content = self.parse_line(line)
if content.strip() == "</s>":
content_type = "Finished"
logger.success("\n[Finished]")
else:
content_type = "Completions"
logger.back(content, end="")
if yield_output:
output = self.message_outputer.output(
content=content, content_type=content_type
)
yield output
|