Views
No views yet
1
2from peft import PeftModel
3from transformers import PreTrainedTokenizer, PreTrainedModel, AutoTokenizer, AutoModelForCausalLM
4from peft import PeftModelForCausalLM, LoraConfig
5from typing import Optional
6from transformers import GenerationConfig
7import torch
8
9PROMPT_DICT = {
10 "prompt_input": (
11 "Below is a^n instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n"
12 "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
13 ),
14 "prompt_no_input": (
15 "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n"
16 "### Instruction:\n{instruction}\n\n### Response:\n"
17 ),
18}
19
20
21def get_model(model_name_or_path: str, load_in_8bit: bool = True, device_map="auto",
22 cpu: bool = False) -> PreTrainedModel:
23 if cpu:
24 model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map=device_map,
25 low_cpu_mem_usage=True)
26 else:
27 model = AutoModelForCausalLM.from_pretrained(model_name_or_path, load_in_8bit=load_in_8bit,
28 device_map=device_map, torch_dtype=torch.float16)
29
30 return model
31
32
33def get_peft_model(model: PreTrainedModel, lora_model_name_or_path: Optional[str] = None) -> PeftModelForCausalLM:
34 model = PeftModel.from_pretrained(model, lora_model_name_or_path, torch_dtype=torch.float16)
35
36 return model
37
38
39def get_tokenizer(model_name_or_path: str, max_input_len: int) -> PreTrainedTokenizer:
40 tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, model_max_length=max_input_len,
41 padding_side="right", use_fast=False)
42
43 return tokenizer
44
45
46def get_llm_inference_model(base_model_name_or_path: str, lora_model_name_or_path: str, load_in_8bit: bool,
47 device_map) -> PeftModel:
48 cpu = True if not torch.cuda.is_available() else False
49
50 model = get_model(base_model_name_or_path, load_in_8bit, device_map, cpu=cpu)
51
52 model = get_peft_model(model, lora_model_name_or_path=lora_model_name_or_path)
53
54 if not load_in_8bit:
55 model.half()
56
57 model.eval()
58
59 if torch.__version__ >= "2":
60 model = torch.compile(model)
61
62 return model
63
64
65def generate_prompt(example):
66 return (
67 PROMPT_DICT["prompt_input"].format_map(example)
68 if example["input"]
69 else PROMPT_DICT["prompt_no_input"].format_map(example)
70 )
71
72
73def infer(instruction: str, input_text: Optional[str] = None, temperature: float = 0.1, top_p: float = 0.95,
74 max_new_tokens: int = 512, early_stopping: bool = True, do_sample: bool = True,
75 repetition_penalty: float = 2.5) -> str:
76 prompt = generate_prompt({"instruction": instruction, "input": input_text})
77
78 tokenized_inputs = tokenizer(prompt, return_tensors="pt")
79
80 device = "cuda" if torch.cuda.is_available() else "cpu"
81
82 input_ids = tokenized_inputs["input_ids"].to(device)
83
84 generation_config = GenerationConfig(temperature=temperature, top_p=top_p, do_sample=do_sample,
85 repetition_penalty=repetition_penalty, early_stopping=early_stopping)
86
87 with torch.inference_mode():
88 generation_output = model.generate(input_ids=input_ids, generation_config=generation_config,
89 return_dict_in_generate=True, max_new_tokens=max_new_tokens)
90
91 output = generation_output.sequences[0]
92
93 output = tokenizer.decode(output, skip_special_tokens=True)
94
95 return output.split("### Response:")[1].strip()
96
97
98base_model_name_or_path = "bigscience/bloom-560m"
99
100lora_model_name_or_path = "crayon-coe/alpaca-bloom-560m-en"
101
102model = get_llm_inference_model(base_model_name_or_path, lora_model_name_or_path, True, "auto")
103
104tokenizer = get_tokenizer(base_model_name_or_path, 512)
105
106context = "Write a letter expressing your love for computers"
107
108output = infer(context)
109
110print(output)
111
112# Output
113# I am so grateful to have been able access this wonderful computer system and its amazing features, which I can now use daily with ease.
114#
115# My heartfelt thanks go out in advance of all my friends who are using it as well.
116# Thank you again!
1171{
2 "max_input_len": 512,
3 "load_in_8bit": True,
4 "model_name_or_path": "bigscience/bloom-560m",
5 "device_map": "auto",
6 "bias": "none",
7 "lora_dropout": 0.05,
8 "lora_alpha": 32,
9 "target_modules": ["query_key_value"],
10 "task_type": "CAUSAL_LM",
11 "lora_r": 16,
12 "pad_to_multiple_of": 8,
13 "num_train_epochs": 3,
14 "learning_rate": 0.0003,
15 "gradient_accumulation_steps": 16,
16 "per_device_train_batch_size": 8,
17 "val_set_size": 500,
18 "save_steps": 200,
19 "eval_steps": 200,
20 "evaluation_strategy": "steps",
21 "save_strategy": "steps"
22}1# coding=utf-8
2# Code 99.99% copied and adapted from:
3# https://github.com/bofenghuang/vigogne
4# https://colab.research.google.com/drive/1jCkpikz0J2o20FBQmYmAGdiKmJGOMo-o?usp=sharing#scrollTo=DpYr24pR8T_0
5
6
7import os
8import sys
9from dataclasses import dataclass
10from typing import Dict, List, Optional, Sequence
11
12import bitsandbytes as bnb
13import fire
14import torch
15import transformers
16from datasets import load_dataset
17from peft import LoraConfig, TaskType, get_peft_model, get_peft_model_state_dict, prepare_model_for_int8_training
18from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaTokenizer
19
20IGNORE_INDEX = -100
21DEFAULT_PAD_TOKEN = "[PAD]"
22
23PROMPT_DICT = {
24 "prompt_input": (
25 "Below is a^n instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n"
26 "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
27 ),
28 "prompt_no_input": (
29 "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n"
30 "### Instruction:\n{instruction}\n\n### Response:\n"
31 ),
32}
33
34
35def generate_prompt(example):
36 return (
37 PROMPT_DICT["prompt_input"].format_map(example)
38 if example["input"]
39 else PROMPT_DICT["prompt_no_input"].format_map(example)
40 )
41
42
43# Modified from: https://github.com/bofenghuang/stanford_alpaca/blob/eb5b171d9b103a12a8e14e0edca9cbc45fe1d512/train.py#L166-L182
44# Almost same to transformers.DataCollatorForSeq2Seq
45@dataclass
46class DataCollatorForSupervisedDataset(object):
47 """Collate examples for supervised fine-tuning."""
48
49 tokenizer: transformers.PreTrainedTokenizer
50 pad_to_multiple_of: Optional[int] = None
51
52 def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
53 # dtype = torch.long
54 # input_ids, labels = tuple([torch.LongTensor(instance[key]) for instance in instances] for key in ("input_ids", "labels"))
55 input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
56
57 if self.pad_to_multiple_of is not None:
58 max_length_index, max_length = max(enumerate([len(input_ids_) for input_ids_ in input_ids]),
59 key=lambda x: x[1])
60 # int(math.ceil
61 n_padding = ((max_length // self.pad_to_multiple_of) + 1) * self.pad_to_multiple_of - max_length
62 # Pad the longest example to pad_to_multiple_of * N
63 input_ids[max_length_index].extend([self.tokenizer.pad_token_id] * n_padding)
64 labels[max_length_index].extend([IGNORE_INDEX] * n_padding)
65
66 input_ids = [torch.LongTensor(input_ids_) for input_ids_ in input_ids]
67 labels = [torch.LongTensor(labels_) for labels_ in labels]
68
69 input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True,
70 padding_value=self.tokenizer.pad_token_id)
71 labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
72
73 return dict(input_ids=input_ids, labels=labels, attention_mask=input_ids.ne(self.tokenizer.pad_token_id))
74
75
76def train(model_name_or_path: str, output_dir: str, data_path: str, val_set_size: int = 500,
77 model_max_length: int = 512, lora_r: int = 16, lora_alpha: int = 32, lora_dropout: float = 0.05,
78 target_modules: List[str] = ["query_key_value"], num_train_epochs: int = 3, learning_rate: float = 0.0001,
79 per_device_train_batch_size: int = 8, gradient_accumulation_steps: int = 16, **kwargs):
80 device_map = "auto"
81
82 model = AutoModelForCausalLM.from_pretrained(model_name_or_path, load_in_8bit=True, device_map=device_map)
83
84 tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, model_max_length=model_max_length,
85 padding_side="right", use_fast=False)
86
87 model = prepare_model_for_int8_training(model)
88
89 lora_config = LoraConfig(r=lora_r, lora_alpha=lora_alpha, target_modules=target_modules, lora_dropout=lora_dropout,
90 bias="none", task_type=TaskType.CAUSAL_LM)
91
92 model = get_peft_model(model, lora_config)
93
94 model.print_trainable_parameters()
95
96 # Load data
97 data = load_dataset("json", data_files=data_path)
98
99 def preprocess_function(example):
100 # Format prompt
101 user_prompt = generate_prompt(example)
102
103 # Get prompt length for masking
104 len_user_prompt_tokens = len(tokenizer(user_prompt, truncation=True)["input_ids"])
105
106 input_ids = tokenizer(user_prompt + example["output"] + tokenizer.eos_token, truncation=True)["input_ids"]
107 labels = [IGNORE_INDEX] * len_user_prompt_tokens + input_ids[len_user_prompt_tokens:]
108
109 return {"input_ids": input_ids, "labels": labels}
110
111 if val_set_size > 0:
112 train_val = data["train"].train_test_split(test_size=val_set_size, shuffle=True, seed=42)
113 train_data = train_val["train"].shuffle().map(preprocess_function, remove_columns=data["train"].column_names)
114 val_data = train_val["test"].map(preprocess_function, remove_columns=data["train"].column_names)
115 else:
116 train_data = data["train"].shuffle().map(preprocess_function, remove_columns=data["train"].column_names)
117 val_data = None
118
119 trainer = transformers.Trainer(
120 model=model,
121 train_dataset=train_data,
122 eval_dataset=val_data,
123 args=transformers.TrainingArguments(
124 per_device_train_batch_size=per_device_train_batch_size,
125 gradient_accumulation_steps=gradient_accumulation_steps,
126 num_train_epochs=num_train_epochs,
127 learning_rate=learning_rate,
128 fp16=True,
129 output_dir=output_dir,
130 load_best_model_at_end=True if val_set_size > 0 else False,
131 **kwargs,
132 ),
133 data_collator=DataCollatorForSupervisedDataset(tokenizer=tokenizer, pad_to_multiple_of=8),
134 )
135 print(trainer.args)
136
137 # Silence the warnings. Please re-enable for inference!
138 model.config.use_cache = False
139
140 old_state_dict = model.state_dict
141 model.state_dict = (lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict())).__get__(model,
142 type(model))
143
144 if torch.__version__ >= "2" and sys.platform != "win32":
145 model = torch.compile(model)
146
147 trainer.train()
148
149 model.save_pretrained(output_dir)
150
151
152if __name__ == "__main__":
153 fire.Fire(train)
154