Spaces:
Runtime error
Runtime error
File size: 2,755 Bytes
476ac07 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
# Copyright (c) OpenMMLab. All rights reserved.
import torch
from peft import LoraConfig
from transformers import (AutoModelForCausalLM, AutoTokenizer,
BitsAndBytesConfig)
from xtuner.model import SupervisedFinetune
__all__ = ['build_model', 'build_lora_model', 'build_qlora_model']
def build_qlora_model(model_name_or_path,
quantization_config=None,
lora_config=None,
return_tokenizer=True):
if quantization_config is None:
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
load_in_8bit=False,
llm_int8_threshold=6.0,
llm_int8_has_fp16_weight=False,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type='nf4')
if lora_config is None:
lora_config = LoraConfig(
r=64,
lora_alpha=16,
lora_dropout=0.1,
bias='none',
task_type='CAUSAL_LM')
llm = AutoModelForCausalLM.from_pretrained(
model_name_or_path,
torch_dtype=torch.float16,
trust_remote_code=True,
quantization_config=quantization_config)
model = SupervisedFinetune(llm, lora=lora_config)
if return_tokenizer:
tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
trust_remote_code=True,
encode_special_tokens=True)
return model.llm, tokenizer
else:
return model.llm
def build_lora_model(model_name_or_path,
lora_config=None,
return_tokenizer=True):
if lora_config is None:
lora_config = LoraConfig(
r=64,
lora_alpha=16,
lora_dropout=0.1,
bias='none',
task_type='CAUSAL_LM')
llm = AutoModelForCausalLM.from_pretrained(
model_name_or_path, torch_dtype=torch.float16, trust_remote_code=True)
model = SupervisedFinetune(llm, lora=lora_config)
if return_tokenizer:
tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
trust_remote_code=True,
encode_special_tokens=True)
return model.llm, tokenizer
else:
return model.llm
def build_model(model_name_or_path, return_tokenizer=True):
model = AutoModelForCausalLM.from_pretrained(
model_name_or_path, torch_dtype=torch.float16, trust_remote_code=True)
if return_tokenizer:
tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
trust_remote_code=True,
encode_special_tokens=True)
return model, tokenizer
else:
return model
|