Spaces:
Running
Running
| import ctranslate2 | |
| from tokenizers import Tokenizer | |
| from huggingface_hub import snapshot_download | |
| import threading | |
| import gradio as gr | |
| from typing import Optional | |
| from queue import Queue | |
| class TokenIteratorStreamer: | |
| def __init__(self, end_token_id: int, timeout: Optional[float] = None): | |
| self.end_token_id = end_token_id | |
| self.queue = Queue() | |
| self.timeout = timeout | |
| def put(self, token: int): | |
| self.queue.put(token, timeout=self.timeout) | |
| def __iter__(self): | |
| return self | |
| def __next__(self): | |
| token = self.queue.get(timeout=self.timeout) | |
| if token == self.end_token_id: | |
| raise StopIteration() | |
| else: | |
| return token | |
| def generate_prompt(history): | |
| prompt = f'[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n' | |
| for chain in history[-2:-1]: | |
| prompt += f"{chain[0]} [/INST] {chain[1]} </s><s>[INST] " | |
| prompt += f"{history[-1][0]} [/INST]" | |
| tokens = tokenizer.encode(prompt).tokens | |
| return tokens | |
| def generate(streamer, history): | |
| def stepResultCallback(result): | |
| streamer.put(result.token_id) | |
| if result.is_last and (result.token_id != end_token_id): | |
| streamer.put(end_token_id) | |
| print(f"step={result.step}, batch_id={result.batch_id}, token={result.token}") | |
| tokens = generate_prompt(history) | |
| results = generator.generate_batch( | |
| [tokens], | |
| include_prompt_in_result=False, | |
| max_length = 256, | |
| callback = stepResultCallback | |
| ) | |
| return results | |
| system_prompt = """\ | |
| You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. | |
| If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\ | |
| """ | |
| model_name = "neongeckocom/Llama-2-7b-chat-hf" | |
| model_path = snapshot_download(repo_id=model_name) | |
| generator = ctranslate2.Generator(model_path, intra_threads=2) | |
| tokenizer = Tokenizer.from_pretrained(model_name) | |
| end_token = "</s>" | |
| end_token_id = tokenizer.encode(end_token).tokens[0] | |
| with gr.Blocks() as demo: | |
| chatbot = gr.Chatbot() | |
| msg = gr.Textbox() | |
| clear = gr.Button("Clear") | |
| def user(user_message, history): | |
| return "", history + [[user_message, None]] | |
| def bot(history): | |
| bot_message_tokens = [] | |
| streamer = TokenIteratorStreamer(end_token_id = end_token_id) | |
| generation_thread = threading.Thread(target=generate, args=(streamer, history)) | |
| generation_thread.start() | |
| history[-1][1] = "" | |
| for token in streamer: | |
| bot_message_tokens.append(token) | |
| history[-1][1] = tokenizer.decode(bot_message_tokens) | |
| yield history | |
| generation_thread.join() | |
| msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then( | |
| bot, chatbot, chatbot | |
| ) | |
| clear.click(lambda: None, None, chatbot, queue=False) | |
| demo.queue() | |
| if __name__ == "__main__": | |
| demo.launch() |