| import gradio as gr | |
| import hf_transfer | |
| from transformers import AutoModelForCausalLM, AutoTokenizer,StoppingCriteriaList,TextIteratorStreamer | |
| from threading import Thread | |
| import os | |
| HFTOKEN=os.getenv("hftoken") | |
| model = AutoModelForCausalLM.from_pretrained( | |
| "kubernetes-bad/chargen-v2", | |
| token = HFTOKEN | |
| ) | |
| tknz=AutoTokenizer.from_pretrained("kubernetes-bad/chargen-v2",token=HFTOKEN) | |
| """ | |
| For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference | |
| """ | |
| def respond( | |
| message, | |
| history: list[tuple[str, str]], | |
| system_message, | |
| max_tokens, | |
| temperature, | |
| top_p, | |
| ): | |
| messages = [{"role": "system", "content": system_message}] | |
| for val in history: | |
| if val[0]: | |
| messages.append({"role": "user", "content": val[0]}) | |
| if val[1]: | |
| messages.append({"role": "assistant", "content": val[1]}) | |
| messages.append({"role": "user", "content": message}) | |
| response = "" | |
| model_inputs = tokenizer.build_chat_input(history=messages, role='user').input_ids.to( | |
| next(model.parameters()).device) | |
| streamer = TextIteratorStreamer(tokenizer, timeout=600, skip_prompt=True) | |
| eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"), | |
| tokenizer.get_command("<|observation|>")] | |
| generate_kwargs = { | |
| "input_ids": model_inputs, | |
| "streamer": streamer, | |
| "max_new_tokens": max_tokens, | |
| "do_sample": True, | |
| "top_p": top_p, | |
| "temperature": temperature, | |
| "stopping_criteria": StoppingCriteriaList([stop]), | |
| "repetition_penalty": 1, | |
| "eos_token_id": eos_token_id, | |
| } | |
| t = Thread(target=model.generate, kwargs=generate_kwargs) | |
| for new_token in streamer: | |
| if new_token and '<|user|>' in new_token: | |
| new_token = new_token.split('<|user|>')[0] | |
| if new_token: | |
| history[-1][1] += new_token | |
| yield history | |
| """ | |
| For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface | |
| """ | |
| js_func = """ | |
| function refresh() { | |
| const url = new URL(window.location); | |
| if (url.searchParams.get('__theme') !== 'dark') { | |
| url.searchParams.set('__theme', 'dark'); | |
| window.location.href = url.href; | |
| } | |
| } | |
| """ | |
| app = gr.ChatInterface( | |
| respond, | |
| js=js_func, | |
| additional_inputs=[ | |
| gr.Textbox(value="You are a bot who generates perfect roleplaying charecters.", label="System message"), | |
| gr.Slider(minimum=1, maximum=4096, value=512, step=1, label="Max new tokens"), | |
| gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"), | |
| gr.Slider( | |
| minimum=0.1, | |
| maximum=1.0, | |
| value=0.95, | |
| step=0.05, | |
| label="Top-p (nucleus sampling)", | |
| ), | |
| ], | |
| ) | |
| if __name__ == "__main__": | |
| app.launch() |