import gradio as gr import os import threading import arrow import time import argparse import logging from dataclasses import dataclass import torch import sentencepiece as spm from transformers import GPTNeoXForCausalLM, GPTNeoXConfig from transformers.generation.streamers import BaseStreamer from huggingface_hub import hf_hub_download, login logger = logging.getLogger() logger.setLevel("INFO") gr_interface = None @dataclass class DefaultArgs: hf_model_name_or_path: str = None spm_model_path: str = None env: str = "dev" port: int = 7860 make_public: bool = False if os.getenv("RUNNING_ON_HF_SPACE"): login(token=os.getenv("HF_TOKEN")) hf_repo = os.getenv("HF_MODEL_REPO") args = DefaultArgs() args.hf_model_name_or_path = hf_repo args.spm_model_path = hf_hub_download(repo_id=hf_repo, filename="sentencepiece.model") else: parser = argparse.ArgumentParser(description="") parser.add_argument("--hf_model_name_or_path", type=str, required=True) parser.add_argument("--spm_model_path", type=str, required=True) parser.add_argument("--env", type=str, default="dev") parser.add_argument("--port", type=int, default=7860) parser.add_argument("--make_public", action='store_true') args = parser.parse_args() def load_model( model_dir, ): config = GPTNeoXConfig.from_pretrained(model_dir) config.is_decoder = True model = GPTNeoXForCausalLM.from_pretrained(model_dir, config=config, torch_dtype=torch.bfloat16) if torch.cuda.is_available(): model = model.to("cuda:0") return model logging.info("Loading model") model = load_model(args.hf_model_name_or_path) sp = spm.SentencePieceProcessor(model_file=args.spm_model_path) logging.info("Finished loading model") class SentencePieceStreamer(BaseStreamer): def __init__(self, sp: spm.SentencePieceProcessor): self.sp = sp self.num_invoked = 0 self.prompt = "" self.generated_text = "" self.ended = False def put(self, t: torch.Tensor): d = t.dim() if d == 1: pass elif d == 2: t = t[0] else: raise NotImplementedError t = [int(x) for x in t.numpy()] text = self.sp.decode_ids(t) if self.num_invoked == 0: self.prompt = text self.num_invoked += 1 return self.generated_text += text logging.debug(f"[streamer]: {self.generated_text}") def end(self): self.ended = True def generate( prompt, max_new_tokens, temperature, repetition_penalty, do_sample, no_repeat_ngram_size, ): log = dict(locals()) logging.debug(log) tokens = sp.encode(prompt) input_ids = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(model.device) streamer = SentencePieceStreamer(sp=sp) max_possilbe_new_tokens = model.config.max_position_embeddings - len(tokens) max_possilbe_new_tokens = min(max_possilbe_new_tokens, max_new_tokens) thr = threading.Thread(target=model.generate, args=(), kwargs=dict( input_ids=input_ids, do_sample=do_sample, temperature=temperature, repetition_penalty=repetition_penalty, no_repeat_ngram_size=no_repeat_ngram_size, max_new_tokens=max_possilbe_new_tokens, streamer=streamer, # max_length=4096, # top_k=100, # top_p=0.9, # num_return_sequences=2, # num_beams=2, )) thr.start() while not streamer.ended: time.sleep(0.05) yield streamer.generated_text # TODO: optimize for final few tokens gen = streamer.generated_text log.update(dict(generation=gen, time=str(arrow.now("+09:00")))) logging.info(log) yield gen def process_feedback( rating, prompt, generation, max_new_tokens, temperature, repetition_penalty, do_sample, no_repeat_ngram_size, ): log = dict(locals()) log["time"] = str(arrow.now("+09:00")) logging.info(log) if gr_interface: gr_interface.close(verbose=False) with gr.Blocks() as gr_interface: with gr.Row(): gr.Markdown("# 日本語 StableLM Pre-Alpha") with gr.Row(): gr.Markdown("Description about this page.　ホゲホゲ") with gr.Row(): # left panel with gr.Column(scale=1): # generation params with gr.Box(): gr.Markdown("パラメータ") # hidden default params do_sample = gr.Checkbox(True, label="Do Sample", visible=False) no_repeat_ngram_size = gr.Slider(0, 10, value=5, step=1, label="No Repeat Ngram Size", visible=False) # visible params max_new_tokens = gr.Slider( 128, min(512, model.config.max_position_embeddings), value=128, step=128, label="max tokens", info="生成するトークンの最大数を指定する", ) temperature = gr.Slider( 0, 1, value=0.7, step=0.05, label="temperature", info="低い値は出力をより集中させて決定論的にする") repetition_penalty = gr.Slider( 1, 1.5, value=1.2, step=0.05, label="frequency penalty", info="高い値はAIが繰り返す可能性を減少させる") # grouping params for easier reference gr_params = [ max_new_tokens, temperature, repetition_penalty, do_sample, no_repeat_ngram_size, ] # right panel with gr.Column(scale=2): # user input block with gr.Box(): textbox_prompt = gr.Textbox( label="Human", placeholder="AIに続きを書いて欲しいプロンプト", interactive=True, lines=5, value="" ) with gr.Box(): with gr.Row(): btn_submit = gr.Button(value="実行", variant="primary") btn_stop = gr.Button(value="中止", variant="stop") # model output block with gr.Box(): textbox_generation = gr.Textbox( label="AI", lines=5, value="" ) with gr.Box(): with gr.Row(): rating_options = [ "😫すごく悪い", "😞微妙", "😐アリ", "🙂合格", "😄すごく良い", ] btn_ratings = [gr.Button(value=v) for v in rating_options] with gr.Box(): gr.Markdown("TODO：For more feedback link for google form") # event handling inputs = [textbox_prompt] + gr_params click_event = btn_submit.click(generate, inputs, textbox_generation, queue=True) btn_stop.click(None, None, None, cancels=click_event, queue=False) for btn_rating in btn_ratings: btn_rating.click(process_feedback, [btn_rating, textbox_prompt, textbox_generation] + gr_params, queue=False) gr_interface.queue(max_size=32, concurrency_count=2) gr_interface.launch(server_port=args.port, share=args.make_public)