import sys
import torch
from peft import PeftModel, PeftModelForCausalLM, LoraConfig
import transformers
import gradio as gr
import argparse
import warnings
import os
from utils import StreamPeftGenerationMixin,StreamLlamaForCausalLM
assert (
"LlamaTokenizer" in transformers._import_structure["models.llama"]
), "LLaMA is now in HuggingFace's main branch.\nPlease reinstall it: pip uninstall transformers && pip install git+https://github.com/huggingface/transformers.git"
from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig, BitsAndBytesConfig
parser = argparse.ArgumentParser()
parser.add_argument("--model_path", type=str, default="/model/13B_hf")
parser.add_argument("--lora_path", type=str, default="checkpoint-3000")
parser.add_argument("--use_typewriter", type=int, default=1)
parser.add_argument("--use_local", type=int, default=1)
args = parser.parse_args()
print(args)
tokenizer = LlamaTokenizer.from_pretrained(args.model_path)
LOAD_8BIT = True
BASE_MODEL = args.model_path
LORA_WEIGHTS = args.lora_path
# fix the path for local checkpoint
lora_bin_path = os.path.join(args.lora_path, "adapter_model.bin")
print(lora_bin_path)
if not os.path.exists(lora_bin_path) and args.use_local:
pytorch_bin_path = os.path.join(args.lora_path, "pytorch_model.bin")
print(pytorch_bin_path)
if os.path.exists(pytorch_bin_path):
os.rename(pytorch_bin_path, lora_bin_path)
warnings.warn(
"The file name of the lora checkpoint'pytorch_model.bin' is replaced with 'adapter_model.bin'"
)
else:
assert ('Checkpoint is not Found!')
if torch.cuda.is_available():
device = "cuda"
else:
device = "cpu"
try:
if torch.backends.mps.is_available():
device = "mps"
except:
pass
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16
)
if device == "cuda":
model = LlamaForCausalLM.from_pretrained(
BASE_MODEL,
quantization_config=bnb_config,
torch_dtype=torch.float16,
device_map="auto", #{"": 0},
)
model = StreamPeftGenerationMixin.from_pretrained(
model, LORA_WEIGHTS, torch_dtype=torch.float16, device_map="auto", #{"": 0}
)
elif device == "mps":
model = LlamaForCausalLM.from_pretrained(
BASE_MODEL,
device_map={"": device},
torch_dtype=torch.float16,
)
model = StreamPeftGenerationMixin.from_pretrained(
model,
LORA_WEIGHTS,
device_map={"": device},
torch_dtype=torch.float16,
)
else:
model = LlamaForCausalLM.from_pretrained(
BASE_MODEL, device_map={"": device}, low_cpu_mem_usage=True
)
model = StreamPeftGenerationMixin.from_pretrained(
model,
LORA_WEIGHTS,
device_map={"": device},
)
model.config.bos_token_id = tokenizer.bos_token_id = 1
model.config.eos_token_id = tokenizer.eos_token_id = 2
def generate_prompt(instruction, input=None):
if input:
return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
### Instruction:
{instruction}
### Input:
{input}
### Response:"""
else:
return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.
### Instruction:
{instruction}
### Response:"""
if not LOAD_8BIT:
model.half() # seems to fix bugs for some users.
model.eval()
if torch.__version__ >= "2" and sys.platform != "win32":
model = torch.compile(model)
def evaluate(
input,
temperature=0.1,
top_p=0.75,
top_k=40,
num_beams=4,
max_new_tokens=128,
min_new_tokens=1,
repetition_penalty=2.0,
**kwargs,
):
prompt = generate_prompt(input)
inputs = tokenizer(prompt, return_tensors="pt")
input_ids = inputs["input_ids"].to(device)
generation_config = GenerationConfig(
temperature=temperature,
top_p=top_p,
top_k=top_k,
num_beams=num_beams,
bos_token_id=1,
eos_token_id=2,
pad_token_id=0,
max_new_tokens=max_new_tokens, # max_length=max_new_tokens+input_sequence
min_new_tokens=min_new_tokens, # min_length=min_new_tokens+input_sequence
**kwargs,
)
with torch.no_grad():
if args.use_typewriter:
for generation_output in model.stream_generate(
input_ids=input_ids,
generation_config=generation_config,
return_dict_in_generate=True,
output_scores=False,
repetition_penalty=float(repetition_penalty),
):
outputs = tokenizer.batch_decode(generation_output)
show_text = "\n--------------------------------------------\n".join(
[output.split("### Response:")[1].strip().replace('�','')+" ▌" for output in outputs]
)
# if show_text== '':
# yield last_show_text
# else:
yield show_text
yield outputs[0].split("### Response:")[1].strip().replace('�','')
else:
generation_output = model.generate(
input_ids=input_ids,
generation_config=generation_config,
return_dict_in_generate=True,
output_scores=False,
repetition_penalty=1.3,
)
output = generation_output.sequences[0]
output = tokenizer.decode(output).split("### Response:")[1].strip()
print(output)
yield output
gr.Interface(
fn=evaluate,
inputs=[
gr.components.Textbox(
lines=2, label="Input", placeholder="Tell me about alpacas."
),
gr.components.Slider(minimum=0, maximum=1, value=0.1, label="Temperature"),
gr.components.Slider(minimum=0, maximum=1, value=0.75, label="Top p"),
gr.components.Slider(minimum=0, maximum=100, step=1, value=40, label="Top k"),
gr.components.Slider(minimum=1, maximum=10, step=1, value=4, label="Beams Number"),
gr.components.Slider(
minimum=1, maximum=2000, step=1, value=256, label="Max New Tokens"
),
gr.components.Slider(
minimum=1, maximum=300, step=1, value=1, label="Min New Tokens"
),
gr.components.Slider(
minimum=0.1, maximum=10.0, step=0.1, value=2.0, label="Repetition Penalty"
),
],
outputs=[
gr.inputs.Textbox(
lines=25,
label="Output",
)
],
title="Chinese-Vicuna 中文小羊驼",
description="中文小羊驼由各种高质量的开源instruction数据集，结合Alpaca-lora的代码训练而来，模型基于开源的llama7B，主要贡献是对应的lora模型。由于代码训练资源要求较小，希望为llama中文lora社区做一份贡献。",
).queue().launch(share=True)