import json import sys from threading import Thread from queue import Queue import os import gradio as gr import torch from peft import PeftModel from transformers import GenerationConfig, AutoTokenizer, AutoModelForCausalLM import time os.environ['MallocStackLogging'] = '0' if torch.cuda.is_available(): device = "auto" else: device = "cpu" def reformat_sft(instruction, input): if input: prefix = ( "Below is an instruction that describes a task, paired with an input that provides further context. " "Write a response that appropriately completes the request.\n" f"### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:" ) else: prefix = ( "Below is an instruction that describes a task. " "Write a response that appropriately completes the request.\n" f"### Instruction:\n{instruction}\n\n### Response:" ) return prefix class TextIterStreamer: def __init__(self, tokenizer, skip_prompt=True, skip_special_tokens=False): self.tokenizer = tokenizer self.skip_prompt = skip_prompt self.skip_special_tokens = skip_special_tokens self.tokens = [] self.text_queue = Queue() # self.text_queue = [] self.next_tokens_are_prompt = True def put(self, value): if self.skip_prompt and self.next_tokens_are_prompt: self.next_tokens_are_prompt = False else: if len(value.shape) > 1: value = value[0] self.tokens.extend(value.tolist()) word = self.tokenizer.decode(self.tokens, skip_special_tokens=self.skip_special_tokens) # self.text_queue.append(word) self.text_queue.put(word) def end(self): # self.text_queue.append(None) self.text_queue.put(None) def __iter__(self): return self def __next__(self): value = self.text_queue.get() if value is None: raise StopIteration() else: return value def main( base_model: str = "", lora_weights: str = "", share_gradio: bool = False, ): tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( base_model, device_map=device, trust_remote_code=True, torch_dtype=torch.float16 ) if lora_weights: model = PeftModel.from_pretrained( model, lora_weights ) model.eval() def evaluate( instruction, temperature=0.1, top_p=0.75, max_new_tokens=128, repetition_penalty=1.1, **kwargs, ): print(instruction, temperature, top_p, max_new_tokens, repetition_penalty, **kwargs) if not instruction: return prompt = reformat_sft(instruction, "") inputs = tokenizer(prompt, return_tensors="pt") input_ids = inputs["input_ids"].cuda() if not (1 > temperature > 0): temperature = 1 if not (1 > top_p > 0): top_p = 1 if not (2000 > max_new_tokens > 0): max_new_tokens = 200 if not (5 > repetition_penalty > 0): repetition_penalty = 1.1 output = ['', '', ''] for i in range(3): if i > 0: time.sleep(0.5) streamer = TextIterStreamer(tokenizer) generation_config = dict( temperature=temperature, top_p=top_p, max_new_tokens=max_new_tokens, do_sample=True, repetition_penalty=repetition_penalty, streamer=streamer, ) c = Thread(target=lambda: model.generate(input_ids=input_ids, **generation_config)) c.start() for text in streamer: output[i] = text yield output[0], output[1], output[2] print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())) print(output) def fk_select(select_option): def inner(context, answer1, answer2, answer3, fankui): print("反馈", select_option, context, answer1, answer2, answer3, fankui) gr.Info("反馈成功") data = { "context": context, "answer": [answer1, answer2, answer3], "choose": "" } if select_option == 1: data["choose"] = answer1 elif select_option == 2: data["choose"] = answer2 elif select_option == 3: data["choose"] = answer3 elif select_option == 4: data["choose"] = fankui with open("fankui.jsonl", 'a+', encoding="utf-8") as f: f.write(json.dumps(data, ensure_ascii=False) + "\n") return inner with gr.Blocks() as demo: gr.Markdown( "# 云起无垠SecGPT模型RLHF测试\n\nHuggingface: https://huggingface.co/w8ay/secgpt\nGithub: https://github.com/Clouditera/secgpt") with gr.Row(): with gr.Column(): # 列排列 context = gr.Textbox( lines=3, label="Instruction", placeholder="Tell me ..", ) temperature = gr.Slider( minimum=0, maximum=1, value=0.3, label="Temperature" ) topp = gr.Slider( minimum=0, maximum=1, value=0.7, label="Top p" ) max_tokens = gr.Slider( minimum=1, maximum=2000, step=1, value=300, label="Max tokens" ) repetion = gr.Slider( minimum=0, maximum=10, value=1.1, label="repetition_penalty" ) with gr.Column(): answer1 = gr.Textbox( lines=4, label="回答1", ) fk1 = gr.Button("选这个") answer2 = gr.Textbox( lines=4, label="回答2", ) fk2 = gr.Button("选这个") answer3 = gr.Textbox( lines=4, label="回答3", ) fk3 = gr.Button("选这个") fankui = gr.Textbox( lines=4, label="反馈回答", ) fk4 = gr.Button("都不好，反馈") with gr.Row(): submit = gr.Button("submit", variant="primary") gr.ClearButton([context, answer1, answer2, answer3, fankui]) submit.click(fn=evaluate, inputs=[context, temperature, topp, max_tokens, repetion], outputs=[answer1, answer2, answer3]) fk1.click(fn=fk_select(1), inputs=[context, answer1, answer2, answer3, fankui]) fk2.click(fn=fk_select(2), inputs=[context, answer1, answer2, answer3, fankui]) fk3.click(fn=fk_select(3), inputs=[context, answer1, answer2, answer3, fankui]) fk4.click(fn=fk_select(4), inputs=[context, answer1, answer2, answer3, fankui]) demo.queue().launch(server_name="0.0.0.0", share=True) # Old testing code follows. if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description='云起无垠SecGPT模型RLHF测试') parser.add_argument("--base_model", type=str, required=True, help="基础模型") parser.add_argument("--lora", type=str, help="lora模型") parser.add_argument("--share_gradio", type=bool, default=False, help="开放外网访问") args = parser.parse_args() main(args.base_model, args.lora, args.share_gradio)