Spaces:
Runtime error
Runtime error
File size: 2,528 Bytes
9131d62 | 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 | from peft import prepare_model_for_kbit_training
from peft import LoraConfig, get_peft_model
import transformers
'''Dataset loading'''
from datasets import load_dataset
def process_dataset(sample, max_target_length: int = 4000):
'''Process dataset labels if necessary'''
model_inputs = tokenizer(
sample["summary"],
truncation=True,
)
labels = tokenizer(
sample["opinion"], max_length=max_target_length, truncation=True
)
model_inputs["labels"] = labels["input_ids"]
return model_inputs
def load_case_dataset(dataset_path:str=''):
'''dataset loader function, needs dataset path'''
dataset = load_dataset('json', data_files='chunked_case_data.json', field='data')
dataset = dataset.map(process_dataset, batched=True)
dataset = dataset.map(lambda samples: tokenizer(samples["opinion"]), batched=True)
return dataset
def print_trainable_parameters(model):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
all_param += param.numel()
if param.requires_grad:
trainable_params += param.numel()
print(
f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
)
def train_model(model, save_path:str=''):
model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=8,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
)
tokenizer.pad_token = tokenizer.eos_token
trainer = transformers.Trainer(
model=model,
train_dataset=dataset['train'],
args=transformers.TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=100,
max_steps=200,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir='outputs'
),
data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # silence the warnings. Please re-enable for inference!
trainer.train()
model = get_peft_model(model, config)
print_trainable_parameters(model)
if save_path:
model.save_pretrained('./llama2_lora_4k/')
return model
|