Spaces:
Sleeping
Sleeping
import os | |
from typing import Iterator, List, Tuple | |
from text_generation import Client | |
model_id = "mistralai/Mistral-7B-Instruct-v0.1" | |
API_URL = "https://api-inference.huggingface.co/models/" + model_id | |
HF_TOKEN = os.environ.get("HF_READ_TOKEN", None) | |
client = Client( | |
API_URL, | |
headers={"Authorization": f"Bearer {HF_TOKEN}"}, | |
) | |
EOS_STRING = "</s>" | |
EOT_STRING = "<EOT>" | |
def _get_prompt( | |
message: str, chat_history: List[Tuple[str, str]], system_prompt: str | |
) -> str: | |
""" | |
Get the prompt to send to the model. | |
:param message: The message to send to the model. | |
:param chat_history: The chat history. | |
:param system_prompt: The system prompt. | |
:return: The prompt to send to the model. | |
""" | |
texts = [f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"] | |
do_strip = False | |
for user_input, response in chat_history: | |
user_input = user_input.strip() if do_strip else user_input | |
do_strip = True | |
texts.append(f"{user_input} [/INST] {response.strip()} </s><s>[INST] ") | |
message = message.strip() if do_strip else message | |
texts.append(f"{message} [/INST]") | |
return "".join(texts) | |
def run( | |
message: str, | |
chat_history: List[Tuple[str, str]], | |
system_prompt: str, | |
max_new_tokens: int = 2048, | |
temperature: float = 0.1, | |
top_p: float = 0.9, | |
top_k: int = 50, | |
) -> Iterator[str]: | |
""" | |
Run the model. | |
:param message: The message to send to the model. | |
:param chat_history: The chat history. | |
:param system_prompt: The system prompt. | |
:param max_new_tokens: The maximum number of tokens to generate. | |
:param temperature: The temperature. | |
:param top_p: The top p. | |
:param top_k: The top k. | |
:return: The generated text. | |
""" | |
prompt = _get_prompt(message, chat_history, system_prompt) | |
generate_kwargs = dict( | |
max_new_tokens=max_new_tokens, | |
do_sample=True, | |
top_p=top_p, | |
top_k=top_k, | |
temperature=temperature, | |
) | |
stream = client.generate_stream(prompt, **generate_kwargs) | |
output = "" | |
for response in stream: | |
if any( | |
[end_token in response.token.text for end_token | |
in [EOS_STRING, EOT_STRING]] | |
): | |
return output | |
else: | |
output += response.token.text | |
yield output | |
return output | |