from threading import Thread import torch import gradio as gr from transformers import pipeline,AutoTokenizer, AutoModelForCausalLM from peft import PeftModel, PeftConfig import re torch_device = "cuda" if torch.cuda.is_available() else "cpu" print("Running on device:", torch_device) print("CPU threads:", torch.get_num_threads()) peft_model_id = "ldhldh/1.3_40kstep" #peft_model_id = "ldhldh/polyglot-ko-1.3b_lora_big_tern_30kstep" # 20k or 30k #18k > 상대의 말까지 하는 이슈가 있음 #8k > 약간 아쉬운가? config = PeftConfig.from_pretrained(peft_model_id) base_model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path) tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path) #base_model = AutoModelForCausalLM.from_pretrained("EleutherAI/polyglot-ko-3.8b") #tokenizer = AutoTokenizer.from_pretrained("EleutherAI/polyglot-ko-3.8b") base_model.eval() base_model.config.use_cache = True model = PeftModel.from_pretrained(base_model, peft_model_id, device_map="auto") model.eval() model.config.use_cache = True def gen(x, top_p, top_k, temperature, max_new_tokens, repetition_penalty): gened = model.generate( **tokenizer( f"{x}", return_tensors='pt', return_token_type_ids=False ), #bad_words_ids = bad_words_ids , max_new_tokens=max_new_tokens, min_new_tokens = 5, exponential_decay_length_penalty = (max_new_tokens/2, 1.1), top_p=top_p, top_k=top_k, temperature = temperature, early_stopping=True, do_sample=True, eos_token_id=2, pad_token_id=2, #stopping_criteria = stopping_criteria, repetition_penalty=repetition_penalty, no_repeat_ngram_size = 2 ) model_output = tokenizer.decode(gened[0]) return model_output def reset_textbox(): return gr.update(value='') with gr.Blocks() as demo: duplicate_link = "https://huggingface.co/spaces/beomi/KoRWKV-1.5B?duplicate=true" gr.Markdown( "duplicated from beomi/KoRWKV-1.5B, baseModel:EleutherAI/polyglot-ko-1.3b" ) with gr.Row(): with gr.Column(scale=4): user_text = gr.Textbox( placeholder='\\nfriend: 우리 여행 갈래? \\nyou:', label="User input" ) model_output = gr.Textbox(label="Model output", lines=10, interactive=False) button_submit = gr.Button(value="Submit") with gr.Column(scale=1): max_new_tokens = gr.Slider( minimum=1, maximum=200, value=20, step=1, interactive=True, label="Max New Tokens", ) top_p = gr.Slider( minimum=0.05, maximum=1.0, value=0.8, step=0.05, interactive=True, label="Top-p (nucleus sampling)", ) top_k = gr.Slider( minimum=5, maximum=100, value=30, step=5, interactive=True, label="Top-k (nucleus sampling)", ) temperature = gr.Slider( minimum=0.1, maximum=2.0, value=0.5, step=0.1, interactive=True, label="Temperature", ) repetition_penalty = gr.Slider( minimum=1.0, maximum=3.0, value=1.2, step=0.1, interactive=True, label="repetition_penalty", ) button_submit.click(gen, [user_text, top_p, top_k, temperature, max_new_tokens, repetition_penalty], model_output) demo.queue(max_size=32).launch(enable_queue=True)