Spaces:
Runtime error
Runtime error
File size: 7,062 Bytes
b11f272 |
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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 |
import json
import random
import re
from collections import defaultdict
from typing import Any, List, Dict, Callable, Union, Optional
# from vllm import LLM, SamplingParams
import regex
import numpy as np
from huggingface_hub import InferenceClient
from tqdm import tqdm
# from config import *
ROLE_SYSTEM = 'system'
ROLE_USER = 'user'
ROLE_ASSISTANT = 'assistant'
CHAT_FORMATS = {
"mistralai": "<s>[INST] {prompt} [/INST]",
"openchat": "GPT4 User: {prompt}<|end_of_turn|>GPT4 Assistant:",
"meta-llama": """[INST] <<SYS>>
You answer questions directly.
<</SYS>>
{prompt}[/INST]""",
"mosaicml": """<|im_start|>system
A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.<|im_end|>
<|im_start|>user
{prompt}<|im_end|>
<|im_start|>assistant""",
"lmsys": "USER: {prompt}\nASSISTANT:",
}
LLAMA_TEMPLATE = """<s>[INST] <<SYS>>
{system_prompt}
<</SYS>>
{user_message} [/INST]"""
MISTRAL_TEMPLATE = """<s>[INST] <<SYS>>
{system_prompt}
<</SYS>> {user_message} [/INST]"""
YI_34B_TEMPLATE = """<|im_start|>system
{system_prompt}<|im_end|>
<|im_start|>user
{user_message}<|im_end|>
<|im_start|>assistant
"""
def extract_json(text: str) -> Dict:
# json_string_match = re.search(r"json\s+(.+?)\s+", text, re.DOTALL)
# Assume it's goind to be like: "Guess": "A" or "Guess": "B"
# Now it's either true or false
# print(text)
text = text.replace('\\', '\\\\')
try:
rslt = json.loads(text)
except Exception as e:
# print(e)
# print(text)
rslt = None
return rslt
def mixtral_prompt_formatter(messages: List[Dict[str, str]]) -> str:
"""
<s>[INST] <<SYS>>
{system_prompt}
<</SYS>>
{user_prompt} [/INST]
"""
assert len(messages) >= 2 # must be at least a system and a user
r = f'<s>[INST] <<SYS>>\n{messages[0]["content"]}\n<</SYS>>\n{messages[1]["content"]} [/INST]'
for msg in messages[2:]:
role, content = msg['role'], msg['content']
if role == ROLE_SYSTEM:
assert ValueError
elif role == ROLE_USER:
if r.endswith('</s>'):
r += '<s>'
r += f'[INST] {content} [/INST]'
elif role == ROLE_ASSISTANT:
r += f'{content}</s>'
else:
raise ValueError
return r
def llama_prompt_formatter(messages: List[Dict[str, str]]) -> str:
"""
<s>[INST] <<SYS>>
{system_prompt}
<</SYS>>
{user_message} [/INST]
"""
assert len(messages) >= 2 # must be at least a system and a user
r = f'<s>[INST] <<SYS>>\n{messages[0]["content"]}\n<</SYS>>\n\n{messages[1]["content"]} [/INST]'
for msg in messages[2:]:
role, content = msg['role'], msg['content']
if role == ROLE_SYSTEM:
assert ValueError
elif role == ROLE_USER:
if r.endswith('</s>'):
r += '<s>'
r += f'[INST] {content} [/INST]'
elif role == ROLE_ASSISTANT:
r += f'{content}</s>'
else:
raise ValueError
return r
def yi_prompt_formatter(messages: List[Dict[str, str]]) -> str:
"""
<|im_start|>system
{system_prompt}<|im_end|>
<|im_start|>user
{user_message}<|im_end|>
<|im_start|>assistant
"""
assert len(messages) >= 2 # must be at least a system and a user
r = f'<|im_start|>system\n{messages[0]["content"]}<|im_end|>\n<|im_start|>user\n{messages[1]["content"]}<|im_end|>\n'
for i in range(2, len(messages)):
msg = messages[i]
role, content = msg['role'], msg['content']
if role == ROLE_SYSTEM:
assert ValueError
elif role == ROLE_USER:
r += f'<|im_start|>user\n{content}<|im_end|>\n'
if i == len(messages) - 1:
r += '<|im_start|>assistant\n'
elif role == ROLE_ASSISTANT:
r += f'<|im_start|>assistant\n{content}<|im_end|>\n'
else:
raise ValueError
return r
def find_first_valid_json(s) -> Optional[Dict]:
s = re.sub(r'\\(?!["\\/bfnrt]|u[0-9a-fA-F]{4})', lambda m: m.group(0)[1:], s)
for i in range(len(s)):
if s[i] != '{':
continue
for j in range(i + 1, len(s) + 1):
if s[j - 1] != '}':
continue
try:
potential_json = s[i:j]
json_obj = json.loads(potential_json, strict=False)
return json_obj # Return the first valid JSON object found
except json.JSONDecodeError:
pass # Continue searching if JSON decoding fails
return None # Return None if no valid JSON object is found
class HFAPIModel:
model_name: str
messages: List[Dict[str, str]]
system_prompt: str
formatter: Callable[[List[Dict[str, str]]], str]
def __init__(self, system_prompt: str, model_name: str) -> None:
self.system_prompt = system_prompt
self.model_name = model_name
if 'llama' in model_name:
self.formatter = llama_prompt_formatter
elif 'mistral' in model_name:
self.formatter = mixtral_prompt_formatter
else:
raise NotImplementedError
self.messages = [
{'role': ROLE_SYSTEM, 'content': system_prompt}
]
def __call__(self, user_prompt: str, use_json: bool = False,
temperature: float = 0, timeout: float = None,
cache: bool = True) -> Union[str, Dict]:
self.add_message(ROLE_USER, user_prompt)
response = self.get_response(temperature, use_json, timeout, cache)
self.add_message(ROLE_ASSISTANT, response)
return response
def get_response(self, temperature: float, use_json: bool, timeout: float, cache: bool) -> Union[str, Dict]:
"""
Returns the model's response.
If use_json = True, will try its best to return a json dict, but not guaranteed.
If we cannot parse the JSON, we will return the response string directly.
"""
client = InferenceClient(self.model_name, timeout=timeout)
if not cache:
client.headers["x-use-cache"] = "0"
# print(self.formatter(self.messages)) # debug
r = client.text_generation(self.formatter(self.messages),
do_sample=temperature > 0,
temperature=temperature if temperature > 0 else None,
max_new_tokens=512)
if use_json:
obj = find_first_valid_json(r)
if obj is not None:
return obj
return r
def add_message(self, role: str, message: str):
self.messages.append({'role': role, 'content': message})
if __name__ == '__main__':
# model = GPTModel(system_prompt='You are an AI developed by OpenAI.', model_name=GPT_4_MODEL_NAME)
model = HFAPIModel(system_prompt='You are a helpful assistant.', model_name='mistralai/Mixtral-8x7B-Instruct-v0.1')
print(model('Who are you?')) |