Spaces:
Running
Running
File size: 8,333 Bytes
4fe4082 b11f272 487b80b 4fe4082 b11f272 4fe4082 b11f272 487b80b b11f272 4fe4082 487b80b 4fe4082 b11f272 4fe4082 b11f272 4fe4082 b11f272 4fe4082 b11f272 4fe4082 b11f272 4fe4082 b11f272 4fe4082 487b80b 4fe4082 b11f272 4fe4082 b11f272 4fe4082 487b80b 4fe4082 b11f272 4fe4082 b11f272 4fe4082 487b80b |
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 |
from __future__ import annotations
import json
import random
import re
import os
from abc import ABC, abstractmethod
from typing import List, Dict, Union, Optional
from huggingface_hub import InferenceClient
from tenacity import retry, stop_after_attempt, wait_random_exponential
from transformers import AutoTokenizer
# from config import *
ROLE_SYSTEM = 'system'
ROLE_USER = 'user'
ROLE_ASSISTANT = 'assistant'
SUPPORTED_MISTRAL_MODELS = ['mistralai/Mixtral-8x7B-Instruct-v0.1', 'mistralai/Mistral-7B-Instruct-v0.2']
SUPPORTED_NOUS_MODELS = ['NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO',
'NousResearch/Nous-Hermes-2-Mistral-7B-DPO']
SUPPORTED_LLAMA_MODELS = ['meta-llama/Llama-2-70b-chat-hf',
'meta-llama/Llama-2-13b-chat-hf',
'meta-llama/Llama-2-7b-chat-hf']
ALL_SUPPORTED_MODELS = SUPPORTED_MISTRAL_MODELS + SUPPORTED_NOUS_MODELS + SUPPORTED_LLAMA_MODELS
def select_model(model_name: str, system_prompt: str, **kwargs) -> Model:
if model_name in SUPPORTED_MISTRAL_MODELS:
return MistralModel(system_prompt, model_name)
elif model_name in SUPPORTED_NOUS_MODELS:
return NousHermesModel(system_prompt, model_name)
elif model_name in SUPPORTED_LLAMA_MODELS:
return LlamaModel(system_prompt, model_name)
else:
raise ValueError(f'Model {model_name} not supported')
class Model(ABC):
name: str
messages: List[Dict[str, str]]
system_prompt: str
def __init__(self, model_name: str, system_prompt: str):
self.name = model_name
self.system_prompt = system_prompt
self.messages = [
{'role': ROLE_SYSTEM, 'content': system_prompt}
]
@abstractmethod
def __call__(self, *args, **kwargs) -> Union[str, Dict]:
raise NotImplementedError
def add_message(self, role: str, content: str):
assert role in [ROLE_SYSTEM, ROLE_USER, ROLE_ASSISTANT]
self.messages.append({'role': role, 'content': content})
def clear_conversations(self):
self.messages.clear()
self.add_message(ROLE_SYSTEM, self.system_prompt)
def __str__(self) -> str:
return self.name
def __repr__(self) -> str:
return self.name
class HFAPIModel(Model):
def __call__(self, user_prompt: str, *args,
use_json: bool = False,
temperature: float = 0,
timeout: float = None,
cache: bool = False,
json_retry_count: int = 5,
**kwargs) -> 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.
"""
self.add_message(ROLE_USER, user_prompt)
response = self.get_response(temperature, use_json, timeout, cache)
if use_json:
for i in range(json_retry_count):
# cache only if both instruct to do and first try
response = self.get_response(temperature, use_json, timeout, cache and i == 0)
json_obj = self.find_first_valid_json(response)
if json_obj is not None:
response = json_obj
break
self.add_message(ROLE_ASSISTANT, response)
return response
@retry(stop=stop_after_attempt(5), wait=wait_random_exponential(max=10), reraise=True) # retry if exception
def get_response(self, temperature: float, use_json: bool, timeout: float, cache: bool) -> str:
# hf_api_token =
client = InferenceClient(model=self.name, token=os.getenv('HF_API_TOKEN'), timeout=timeout)
# client = InferenceClient(model=self.name, token=random.choice(HF_API_TOKENS), timeout=timeout)
if not cache:
client.headers["x-use-cache"] = "0"
# print(self.formatter(self.messages)) # debug
r = client.text_generation(self.format_messages(),
do_sample=temperature > 0,
temperature=temperature if temperature > 0 else None,
max_new_tokens=4096)
return r
@abstractmethod
def format_messages(self) -> str:
raise NotImplementedError
def get_short_name(self) -> str:
"""
Returns the last part of the model name.
For example, "mistralai/Mixtral-8x7B-Instruct-v0.1" -> "Mixtral-8x7B-Instruct-v0.1"
"""
return self.name.split('/')[-1]
@staticmethod
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) # remove all invalid escapes chars
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 MistralModel(HFAPIModel):
def __init__(self, system_prompt: str, model_name: str = 'mistralai/Mixtral-8x7B-Instruct-v0.1') -> None:
assert model_name in ['mistralai/Mixtral-8x7B-Instruct-v0.1',
'mistralai/Mistral-7B-Instruct-v0.2'], 'Model not supported'
super().__init__(model_name, system_prompt)
def format_messages(self) -> str:
messages = self.messages
# mistral doesn't support system prompt, so we need to convert it to user prompt
if messages[0]['role'] == ROLE_SYSTEM:
assert len(self.messages) >= 2
messages = [{'role' : ROLE_USER,
'content': messages[0]['content'] + '\n' + messages[1]['content']}] + messages[2:]
tokenizer = AutoTokenizer.from_pretrained(self.name)
r = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, max_length=4096)
# print(r)
return r
class NousHermesModel(HFAPIModel):
def __init__(self, system_prompt: str, model_name: str = 'NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO') -> None:
assert model_name in SUPPORTED_NOUS_MODELS, 'Model not supported'
super().__init__(model_name, system_prompt)
def format_messages(self) -> str:
messages = self.messages
assert len(messages) >= 2 # must be at least a system and a user
assert messages[0]['role'] == ROLE_SYSTEM and messages[1]['role'] == ROLE_USER
tokenizer = AutoTokenizer.from_pretrained(self.name)
r = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, max_length=4096)
# print(r)
return r
class LlamaModel(HFAPIModel):
def __init__(self, system_prompt: str, model_name: str = 'meta-llama/Llama-2-70b-chat-hf') -> None:
assert model_name in ['meta-llama/Llama-2-70b-chat-hf',
'meta-llama/Llama-2-13b-chat-hf',
'meta-llama/Llama-2-7b-chat-hf'], 'Model not supported'
super().__init__(model_name, system_prompt)
def format_messages(self) -> str:
"""
<s>[INST] <<SYS>>
{system_prompt}
<</SYS>>
{user_message} [/INST]
"""
messages = self.messages
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 |