Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer | |
| import time | |
| import numpy as np | |
| from torch.nn import functional as F | |
| import os | |
| from threading import Thread | |
| print(f"Starting to load the model to memory") | |
| m = AutoModelForCausalLM.from_pretrained( | |
| "stabilityai/stablelm-tuned-alpha-7b", torch_dtype=torch.float16).cuda() | |
| tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b") | |
| generator = pipeline('text-generation', model=m, tokenizer=tok, device=0) | |
| print(f"Sucessfully loaded the model to the memory") | |
| start_message = """<|SYSTEM|># StableAssistant | |
| - StableAssistant is A helpful and harmless Open Source AI Language Model developed by Stability and CarperAI. | |
| - StableAssistant is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user. | |
| - StableAssistant is more than just an information source, StableAssistant is also able to write poetry, short stories, and make jokes. | |
| - StableAssistant will refuse to participate in anything that could harm a human.""" | |
| class StopOnTokens(StoppingCriteria): | |
| def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: | |
| stop_ids = [50278, 50279, 50277, 1, 0] | |
| for stop_id in stop_ids: | |
| if input_ids[0][-1] == stop_id: | |
| return True | |
| return False | |
| def user(user_message, history): | |
| history = history + [[user_message, ""]] | |
| return "", history, history | |
| def bot(history, curr_system_message): | |
| stop = StopOnTokens() | |
| messages = curr_system_message + \ | |
| "".join(["".join(["<|USER|>"+item[0], "<|ASSISTANT|>"+item[1]]) | |
| for item in history]) | |
| #model_inputs = tok([messages], return_tensors="pt")['input_ids'].cuda()[:, :4096-1024] | |
| model_inputs = tok([messages], return_tensors="pt").to("cuda") | |
| streamer = TextIteratorStreamer(tok, timeout=10., skip_prompt=True, skip_special_tokens=True) | |
| generate_kwargs = dict( | |
| model_inputs, | |
| streamer=streamer, | |
| max_new_tokens=1024, | |
| do_sample=True, | |
| top_p=0.95, | |
| top_k=1000, | |
| temperature=1.0, | |
| num_beams=1, | |
| stopping_criteria=StoppingCriteriaList([stop]) | |
| ) | |
| t = Thread(target=m.generate, kwargs=generate_kwargs) | |
| t.start() | |
| print(history) | |
| for new_text in streamer: | |
| print(new_text) | |
| history[-1][1] += new_text | |
| yield history, history | |
| return history, history | |
| with gr.Blocks() as demo: | |
| history = gr.State([]) | |
| gr.Markdown("## StableLM-Tuned-Alpha-7b Chat") | |
| gr.HTML('''<center><a href="https://huggingface.co/spaces/stabilityai/stablelm-tuned-alpha-chat?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>Duplicate the Space to skip the queue and run in a private space</center>''') | |
| chatbot = gr.Chatbot().style(height=500) | |
| with gr.Row(): | |
| with gr.Column(scale=0.70): | |
| msg = gr.Textbox(label="Chat Message Box", placeholder="Chat Message Box", show_label=False).style(container=False) | |
| with gr.Column(scale=0.30): | |
| with gr.Row(): | |
| submit = gr.Button("Submit") | |
| clear = gr.Button("Clear") | |
| system_msg = gr.Textbox( | |
| start_message, label="System Message", interactive=False, visible=False) | |
| msg.submit(fn=user, inputs=[msg, history], outputs=[msg, chatbot, history], queue=False).then( | |
| fn=bot, inputs=[chatbot, system_msg], outputs=[chatbot, history], queue=True) | |
| submit.click(fn=user, inputs=[msg, history], outputs=[msg, chatbot, history], queue=False).then( | |
| fn=bot, inputs=[chatbot, system_msg], outputs=[chatbot, history], queue=True) | |
| clear.click(lambda: [None, []], None, [chatbot, history], queue=False) | |
| demo.queue(concurrency_count=2) | |
| demo.launch() |