from typing import List, Union from pathlib import Path import gradio as gr import torch import argparse from threading import Thread from transformers import ( AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, ) import warnings import spaces import os warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated') MODEL_PATH = os.environ.get('MODEL_PATH', 'IndexTeam/Index-1.9B-Character') TOKENIZER_PATH = os.environ.get("TOKENIZER_PATH", MODEL_PATH) tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True) def _resolve_path(path: Union[str, Path]) -> Path: return Path(path).expanduser().resolve() @spaces.GPU def hf_gen(dialog: List, top_k, top_p, temperature, repetition_penalty, max_dec_len): """ Generate model output with Huggingface API Args: dialog (List): List of dialog messages. top_k (int): The number of highest probability vocabulary tokens to keep for top-k-filtering. top_p (float): Only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for generation. temperature (float): Strictly positive float value used to modulate the logits distribution. repetition_penalty (float): The parameter for repetition penalty. max_dec_len (int): The maximum numbers of tokens to generate. Yields: str: Real-time generation results of HF model. """ inputs = tokenizer.apply_chat_template(dialog, tokenize=False, add_generation_prompt=False) enc = tokenizer(inputs, return_tensors="pt").to("cuda") streamer = TextIteratorStreamer(tokenizer, **tokenizer.init_kwargs) generation_kwargs = dict( enc, do_sample=True, top_k=int(top_k), top_p=float(top_p), temperature=float(temperature), repetition_penalty=float(repetition_penalty), max_new_tokens=int(max_dec_len), pad_token_id=tokenizer.eos_token_id, streamer=streamer, ) thread = Thread(target=model.generate, kwargs=generation_kwargs) thread.start() answer = "" for new_text in streamer: answer += new_text yield answer[len(inputs):] @spaces.GPU def generate(chat_history: List, query, top_k, top_p, temperature, repetition_penalty, max_dec_len, system_message): """ Generate after hitting "submit" button Args: chat_history (List): List that stores all QA records. query (str): Query of current round. top_p (float): Only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for generation. temperature (float): Strictly positive float value used to modulate the logits distribution. max_dec_len (int): The maximum numbers of tokens to generate. Yields: List: Updated chat_history with the current round QA. """ assert query != "", "Input must not be empty!!!" # apply chat template model_input = [] if system_message: model_input.append({ "role": "system", "content": system_message }) for q, a in chat_history: model_input.append({"role": "user", "content": q}) model_input.append({"role": "assistant", "content": a}) model_input.append({"role": "user", "content": query}) # yield model generation chat_history.append([query, ""]) for answer in hf_gen(model_input, top_k, top_p, temperature, repetition_penalty, max_dec_len): chat_history[-1][1] = answer.strip(tokenizer.eos_token) yield gr.update(value=""), chat_history @spaces.GPU def regenerate(chat_history: List, top_k, top_p, temperature, repetition_penalty, max_dec_len, system_message): """ Re-generate the answer of last round's query Args: chat_history (List): List that stores all QA records. top_p (float): Only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for generation. temperature (float): Strictly positive float value used to modulate the logits distribution. max_dec_len (int): The maximum numbers of tokens to generate. Yields: List: Updated chat_history. """ assert len(chat_history) >= 1, "History is empty. Nothing to regenerate!!" # apply chat template model_input = [] if system_message: model_input.append({ "role": "system", "content": system_message }) for q, a in chat_history[:-1]: model_input.append({"role": "user", "content": q}) model_input.append({"role": "assistant", "content": a}) model_input.append({"role": "user", "content": chat_history[-1][0]}) # yield model generation for answer in hf_gen(model_input, top_k, top_p, temperature, repetition_penalty, max_dec_len): chat_history[-1][1] = answer.strip(tokenizer.eos_token) yield gr.update(value=""), chat_history def clear_history(): """ Clear all chat history Returns: List: Empty chat history """ torch.cuda.empty_cache() return [] def reverse_last_round(chat_history): """ Reverse last round QA and keep the chat history before Args: chat_history (List): List that stores all QA records. Returns: List: Updated chat_history without the last round. """ assert len(chat_history) >= 1, "History is empty. Nothing to reverse!!" return chat_history[:-1] # launch gradio demo with gr.Blocks(theme="soft") as demo: gr.Markdown("""# Index-1.9B-Character Gradio Demo""") with gr.Row(): with gr.Column(scale=1): top_k = gr.Slider(1, 10, value=5, step=1, label="top_k") top_p = gr.Slider(0, 1, value=0.8, step=0.1, label="top_p") temperature = gr.Slider(0.1, 2.0, value=0.3, step=0.1, label="temperature") repetition_penalty = gr.Slider(0.1, 2.0, value=1.1, step=0.1, label="repetition_penalty") max_dec_len = gr.Slider(1, 4096, value=1024, step=1, label="max_dec_len") with gr.Row(): system_message = gr.Textbox(label="System Message", placeholder="Input your system message", value="你是由哔哩哔哩自主研发的大语言模型，名为“Index-1.9B-Character”。你能够根据用户传入的信息，帮助用户完成指定的任务，并生成恰当的、符合要求的回复。") with gr.Column(scale=10): chatbot = gr.Chatbot(bubble_full_width=False, height=500, label='Index-1.9B-Character') user_input = gr.Textbox(label="User", placeholder="Input your query here!", lines=8) with gr.Row(): submit = gr.Button("🚀 Submit") clear = gr.Button("🧹 Clear") regen = gr.Button("🔄 Regenerate") reverse = gr.Button("⬅️ Reverse") submit.click(generate, inputs=[chatbot, user_input, top_k, top_p, temperature, repetition_penalty, max_dec_len, system_message], outputs=[user_input, chatbot]) regen.click(regenerate, inputs=[chatbot, top_k, top_p, temperature, repetition_penalty, max_dec_len, system_message], outputs=[user_input, chatbot]) clear.click(clear_history, inputs=[], outputs=[chatbot]) reverse.click(reverse_last_round, inputs=[chatbot], outputs=[chatbot]) demo.queue().launch()