mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
formatting
This commit is contained in:
51616
parent
404566501b
commit
2666bf2ced
7 changed files with 328 additions and 196 deletions
|
|
@ -1,10 +1,18 @@
|
|||
import transformers
|
||||
from peft import (
|
||||
LoraConfig,
|
||||
)
|
||||
from datasets import load_dataset
|
||||
from modeling_icae_multi_span import ICAE, ModelArguments, DataArguments, TrainingArguments
|
||||
from training_utils import instruct_ft_tokenize_function, DataCollatorForDynamicPadding, train_model
|
||||
from modeling_icae_multi_span import (
|
||||
ICAE,
|
||||
DataArguments,
|
||||
ModelArguments,
|
||||
TrainingArguments,
|
||||
)
|
||||
from peft import LoraConfig
|
||||
from training_utils import (
|
||||
DataCollatorForDynamicPadding,
|
||||
instruct_ft_tokenize_function,
|
||||
train_model,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
|
||||
|
|
@ -12,22 +20,22 @@ def main():
|
|||
|
||||
print(model_args)
|
||||
print(data_args)
|
||||
|
||||
training_args.gradient_checkpointing_kwargs = {"use_reentrant": False} # manually add this argument in the code
|
||||
|
||||
training_args.gradient_checkpointing_kwargs = {
|
||||
"use_reentrant": False
|
||||
} # manually add this argument in the code
|
||||
|
||||
lora_config = LoraConfig(
|
||||
r=model_args.lora_r,
|
||||
lora_alpha=32,
|
||||
lora_dropout=0.05,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM"
|
||||
r=model_args.lora_r, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
|
||||
)
|
||||
|
||||
|
||||
# check model_args.mem_size and min_tokens_for_lm
|
||||
assert (training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)) == 0, "training_args.fixed_mem_size must be a power of 2"
|
||||
|
||||
assert (
|
||||
training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)
|
||||
) == 0, "training_args.fixed_mem_size must be a power of 2"
|
||||
|
||||
memory_size = training_args.fixed_mem_size
|
||||
|
||||
|
||||
train_file = "/path/to/train/file"
|
||||
eval_file = "/path/to/dev/file"
|
||||
|
||||
|
|
@ -41,9 +49,14 @@ def main():
|
|||
model = ICAE(model_args, training_args, lora_config)
|
||||
MEM_TOKENS = list(range(model.vocab_size, model.vocab_size + memory_size))
|
||||
|
||||
train_dataset = train_dataset.map(instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS})
|
||||
eval_dataset = eval_dataset.map(instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS})
|
||||
train_dataset = train_dataset.map(
|
||||
instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS}
|
||||
)
|
||||
eval_dataset = eval_dataset.map(
|
||||
instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS}
|
||||
)
|
||||
|
||||
train_model(model, train_dataset, eval_dataset, training_args)
|
||||
|
||||
main()
|
||||
|
||||
main()
|
||||
|
|
|
|||
|
|
@ -1,11 +1,17 @@
|
|||
import json
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
from transformers import HfArgumentParser, AutoModelForCausalLM
|
||||
from peft import LoraConfig
|
||||
from modeling_icae_multi_span import ICAE, ModelArguments, DataArguments, TrainingArguments
|
||||
import sys
|
||||
|
||||
import torch
|
||||
from modeling_icae_multi_span import (
|
||||
ICAE,
|
||||
DataArguments,
|
||||
ModelArguments,
|
||||
TrainingArguments,
|
||||
)
|
||||
from peft import LoraConfig
|
||||
from safetensors.torch import load_file
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoModelForCausalLM, HfArgumentParser
|
||||
|
||||
# Set the computation device
|
||||
device = "cuda"
|
||||
|
|
@ -16,11 +22,7 @@ model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
|||
|
||||
# Define Lora configuration
|
||||
lora_config = LoraConfig(
|
||||
r=512,
|
||||
lora_alpha=32,
|
||||
lora_dropout=model_args.lora_dropout,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM"
|
||||
r=512, lora_alpha=32, lora_dropout=model_args.lora_dropout, bias="none", task_type="CAUSAL_LM"
|
||||
)
|
||||
|
||||
# Initialize model and send it to CUDA device
|
||||
|
|
@ -29,14 +31,14 @@ model = ICAE(model_args, training_args, lora_config)
|
|||
# Load the fine-tuned checkpoint
|
||||
print(f"Loading trained checkpoint from {training_args.output_dir}")
|
||||
state_dict = load_file(training_args.output_dir)
|
||||
model.load_state_dict(state_dict, strict=False) # only load lora and memory token embeddings
|
||||
model.load_state_dict(state_dict, strict=False) # only load lora and memory token embeddings
|
||||
|
||||
model = model.to(device)
|
||||
|
||||
# Read the data file
|
||||
file_path = "./dev_v2.jsonl"
|
||||
lines = None
|
||||
with open(file_path, "r") as f:
|
||||
with open(file_path) as f:
|
||||
lines = f.readlines()
|
||||
|
||||
# Prepare the model for evaluation
|
||||
|
|
@ -49,24 +51,40 @@ with torch.no_grad():
|
|||
for line in tqdm(lines):
|
||||
# Tokenize input text
|
||||
data = json.loads(line)
|
||||
tokenized_input = model.tokenizer(data['input'], truncation=True, max_length=5120, padding=False, return_attention_mask=False)
|
||||
tokenized_prompt = model.tokenizer(data['prompt'], truncation=False, padding=False, return_attention_mask=False, add_special_tokens=False)
|
||||
tokenized_input = model.tokenizer(
|
||||
data["input"],
|
||||
truncation=True,
|
||||
max_length=5120,
|
||||
padding=False,
|
||||
return_attention_mask=False,
|
||||
)
|
||||
tokenized_prompt = model.tokenizer(
|
||||
data["prompt"],
|
||||
truncation=False,
|
||||
padding=False,
|
||||
return_attention_mask=False,
|
||||
add_special_tokens=False,
|
||||
)
|
||||
# Generate compressed outputs
|
||||
input_ids = torch.LongTensor([tokenized_input['input_ids']]).to(device)
|
||||
input_ids = torch.LongTensor([tokenized_input["input_ids"]]).to(device)
|
||||
memory_slots = model._compress(input_ids)
|
||||
|
||||
|
||||
# decoder input has 3 parts: prefix, memory slots and suffix
|
||||
# the following code is for Mistral tokenizer for example: 733, 16289, 28793 are for the Mistral instruction tempmlate
|
||||
prompt_left_ids = torch.LongTensor([[1, 733, 16289, 28793]]).to(device)
|
||||
prompt_right_ids = [model.ft_token_id] + tokenized_prompt['input_ids'] + [733, 28748, 16289, 28793]
|
||||
prompt_left_ids = torch.LongTensor([[1, 733, 16289, 28793]]).to(device)
|
||||
prompt_right_ids = (
|
||||
[model.ft_token_id] + tokenized_prompt["input_ids"] + [733, 28748, 16289, 28793]
|
||||
)
|
||||
prompt_right_ids = torch.LongTensor([prompt_right_ids]).to(device)
|
||||
|
||||
prompt_left_embs = model.tokens_to_embeddings(prompt_left_ids)
|
||||
prompt_right_embs = model.tokens_to_embeddings(prompt_right_ids)
|
||||
memory_slots = memory_slots.to(prompt_right_embs)
|
||||
|
||||
|
||||
# Concatenate and clone input embeddings
|
||||
decoder_input_embeddings = torch.cat((prompt_left_embs, memory_slots.unsqueeze(0), prompt_right_embs), dim=1)
|
||||
decoder_input_embeddings = torch.cat(
|
||||
(prompt_left_embs, memory_slots.unsqueeze(0), prompt_right_embs), dim=1
|
||||
)
|
||||
output = decoder_input_embeddings.clone()
|
||||
|
||||
generate_text = []
|
||||
|
|
@ -74,29 +92,36 @@ with torch.no_grad():
|
|||
|
||||
# Generate text output
|
||||
for i in range(max_out_length):
|
||||
with model.icae.disable_adapter(): # no independent decoder; use self.icae
|
||||
out = model.icae(inputs_embeds=output, past_key_values=past_key_values, use_cache=True)
|
||||
with model.icae.disable_adapter(): # no independent decoder; use self.icae
|
||||
out = model.icae(
|
||||
inputs_embeds=output, past_key_values=past_key_values, use_cache=True
|
||||
)
|
||||
# out = decoder(inputs_embeds=output, past_key_values=past_key_values, use_cache=True)
|
||||
logit = out.logits[:, -1, :model.vocab_size-1]
|
||||
logit = out.logits[:, -1, : model.vocab_size - 1]
|
||||
past_key_values = out.past_key_values
|
||||
|
||||
next_token_id = torch.argmax(logit, dim=-1)
|
||||
# print(next_token_id)
|
||||
|
||||
if next_token_id.item() == 2: # eos
|
||||
|
||||
if next_token_id.item() == 2: # eos
|
||||
break
|
||||
|
||||
output = model.icae.get_base_model().model.embed_tokens(next_token_id).unsqueeze(1).to(device)
|
||||
output = (
|
||||
model.icae.get_base_model()
|
||||
.model.embed_tokens(next_token_id)
|
||||
.unsqueeze(1)
|
||||
.to(device)
|
||||
)
|
||||
generate_text.append(next_token_id.item())
|
||||
|
||||
generated_text = model.tokenizer.decode(generate_text)
|
||||
|
||||
# Structure output data
|
||||
output_ = {
|
||||
"input": data['input'],
|
||||
"input": data["input"],
|
||||
"prompt": data["prompt"],
|
||||
"output": generated_text,
|
||||
"answer": data["answer"]
|
||||
"answer": data["answer"],
|
||||
}
|
||||
|
||||
f.write(json.dumps(output_) + "\n")
|
||||
|
|
|
|||
|
|
@ -1,10 +1,18 @@
|
|||
import transformers
|
||||
from peft import (
|
||||
LoraConfig,
|
||||
)
|
||||
from datasets import load_dataset
|
||||
from modeling_icae_multi_span import ICAE, ModelArguments, DataArguments, TrainingArguments
|
||||
from training_utils import instruct_ft_tokenize_function, DataCollatorForDynamicPadding, train_model
|
||||
from modeling_icae_multi_span import (
|
||||
ICAE,
|
||||
DataArguments,
|
||||
ModelArguments,
|
||||
TrainingArguments,
|
||||
)
|
||||
from peft import LoraConfig
|
||||
from training_utils import (
|
||||
DataCollatorForDynamicPadding,
|
||||
instruct_ft_tokenize_function,
|
||||
train_model,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
|
||||
|
|
@ -12,22 +20,22 @@ def main():
|
|||
|
||||
print(model_args)
|
||||
print(data_args)
|
||||
|
||||
training_args.gradient_checkpointing_kwargs = {"use_reentrant": False} # manually add this argument in the code
|
||||
|
||||
training_args.gradient_checkpointing_kwargs = {
|
||||
"use_reentrant": False
|
||||
} # manually add this argument in the code
|
||||
|
||||
lora_config = LoraConfig(
|
||||
r=model_args.lora_r,
|
||||
lora_alpha=32,
|
||||
lora_dropout=0.05,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM"
|
||||
r=model_args.lora_r, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
|
||||
)
|
||||
|
||||
|
||||
# check model_args.mem_size and min_tokens_for_lm
|
||||
assert (training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)) == 0, "training_args.fixed_mem_size must be a power of 2"
|
||||
|
||||
assert (
|
||||
training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)
|
||||
) == 0, "training_args.fixed_mem_size must be a power of 2"
|
||||
|
||||
memory_size = training_args.fixed_mem_size
|
||||
|
||||
|
||||
train_file = "/path/to/train/file"
|
||||
eval_file = "/path/to/dev/file"
|
||||
|
||||
|
|
@ -41,9 +49,14 @@ def main():
|
|||
model = ICAE(model_args, training_args, lora_config)
|
||||
MEM_TOKENS = list(range(model.vocab_size, model.vocab_size + memory_size))
|
||||
|
||||
train_dataset = train_dataset.map(instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS})
|
||||
eval_dataset = eval_dataset.map(instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS})
|
||||
train_dataset = train_dataset.map(
|
||||
instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS}
|
||||
)
|
||||
eval_dataset = eval_dataset.map(
|
||||
instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS}
|
||||
)
|
||||
|
||||
train_model(model, train_dataset, eval_dataset, training_args)
|
||||
|
||||
main()
|
||||
|
||||
main()
|
||||
|
|
|
|||
|
|
@ -1,42 +1,43 @@
|
|||
# ICAE that supports multi span concat
|
||||
|
||||
import transformers
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
|
||||
import math
|
||||
import os
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import random
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
from peft import (
|
||||
get_peft_model,
|
||||
)
|
||||
from torch.nn.functional import gelu
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import transformers
|
||||
from peft import get_peft_model
|
||||
from safetensors.torch import load_file
|
||||
from torch.nn.functional import gelu
|
||||
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArguments:
|
||||
model_name_or_path: str = field(default="mistralai/Mistral-7B-v0.1")
|
||||
lora_r: int = field(
|
||||
default=128,
|
||||
metadata={"help": "lora rank"}
|
||||
)
|
||||
lora_dropout: float = field(
|
||||
default=0.05,
|
||||
metadata={"help": "lora dropout"}
|
||||
)
|
||||
lora_r: int = field(default=128, metadata={"help": "lora rank"})
|
||||
lora_dropout: float = field(default=0.05, metadata={"help": "lora dropout"})
|
||||
train: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "if true, the model ckpt will be initialized for training; else, it's for inference"}
|
||||
metadata={
|
||||
"help": "if true, the model ckpt will be initialized for training; else, it's for inference"
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataArguments:
|
||||
data_path: str = field(default=None, metadata={"help": "Path to the training data."})
|
||||
debug_data: bool = field(default=False, metadata={"help": "Enable debug dataset to quickly verify the training process"})
|
||||
debug_data: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Enable debug dataset to quickly verify the training process"},
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainingArguments(transformers.TrainingArguments):
|
||||
|
|
@ -44,7 +45,9 @@ class TrainingArguments(transformers.TrainingArguments):
|
|||
optim: str = field(default="adamw_torch")
|
||||
model_max_length: int = field(
|
||||
default=28000,
|
||||
metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
|
||||
metadata={
|
||||
"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
|
||||
},
|
||||
)
|
||||
fixed_mem_size: int = field(
|
||||
default=128,
|
||||
|
|
@ -68,13 +71,16 @@ class TrainingArguments(transformers.TrainingArguments):
|
|||
)
|
||||
add_special_token_for_lm: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Add a special token for the prompt of language modeling; default: False"},
|
||||
metadata={
|
||||
"help": "Add a special token for the prompt of language modeling; default: False"
|
||||
},
|
||||
)
|
||||
restore_from: str = field(
|
||||
default="",
|
||||
metadata={"help": "The checkpoint that should be restored from for fine-tuning"}
|
||||
metadata={"help": "The checkpoint that should be restored from for fine-tuning"},
|
||||
)
|
||||
|
||||
|
||||
def print_trainable_parameters(model):
|
||||
trainable_parameters = 0
|
||||
all_param = 0
|
||||
|
|
@ -82,7 +88,9 @@ def print_trainable_parameters(model):
|
|||
all_param += param.numel()
|
||||
if param.requires_grad:
|
||||
trainable_parameters += param.numel()
|
||||
print(f"trainable params: {trainable_parameters} || all params: {all_param} || trainable%: {100 * trainable_parameters / all_param}")
|
||||
print(
|
||||
f"trainable params: {trainable_parameters} || all params: {all_param} || trainable%: {100 * trainable_parameters / all_param}"
|
||||
)
|
||||
# for name, param in model.named_parameters():
|
||||
# if param.requires_grad:
|
||||
# print(name, param.shape)
|
||||
|
|
@ -99,46 +107,61 @@ class ICAE(torch.nn.Module):
|
|||
self.model_args = model_args
|
||||
self.training_args = training_args
|
||||
self.model_name = model_args.model_name_or_path
|
||||
self.icae = AutoModelForCausalLM.from_pretrained(self.model_name, torch_dtype=torch.float16 if training_args.bf16 is False else torch.bfloat16, use_flash_attention_2=True, resume_download=True)
|
||||
|
||||
self.training = self.model_args.train
|
||||
|
||||
if self.training: # indepedent model for gradient checkpointing
|
||||
self.decoder = AutoModelForCausalLM.from_pretrained(self.model_name, torch_dtype=torch.float16 if training_args.bf16 is False else torch.bfloat16, use_flash_attention_2=True, resume_download=True)
|
||||
self.icae = AutoModelForCausalLM.from_pretrained(
|
||||
self.model_name,
|
||||
torch_dtype=torch.float16 if training_args.bf16 is False else torch.bfloat16,
|
||||
use_flash_attention_2=True,
|
||||
resume_download=True,
|
||||
)
|
||||
|
||||
self.vocab_size = self.icae.config.vocab_size + 1 # [PAD] token
|
||||
self.training = self.model_args.train
|
||||
|
||||
if self.training: # indepedent model for gradient checkpointing
|
||||
self.decoder = AutoModelForCausalLM.from_pretrained(
|
||||
self.model_name,
|
||||
torch_dtype=torch.float16 if training_args.bf16 is False else torch.bfloat16,
|
||||
use_flash_attention_2=True,
|
||||
resume_download=True,
|
||||
)
|
||||
|
||||
self.vocab_size = self.icae.config.vocab_size + 1 # [PAD] token
|
||||
self.pad_token_id = self.vocab_size - 1
|
||||
self.mean_compression_rate = training_args.mean_compression_rate
|
||||
|
||||
# tunable
|
||||
self.mem_size = self.training_args.fixed_mem_size
|
||||
self.vocab_size_with_mem = self.vocab_size + self.mem_size # so, the mem tokens are in the range [self.vocab_size, self.vocab_size + self.mem_size)
|
||||
self.vocab_size_with_mem = (
|
||||
self.vocab_size + self.mem_size
|
||||
) # so, the mem tokens are in the range [self.vocab_size, self.vocab_size + self.mem_size)
|
||||
|
||||
# special tokens in addition to mem and length tokens
|
||||
self.ae_token_id = self.vocab_size_with_mem + 0
|
||||
self.lm_token_id = self.vocab_size_with_mem + 1
|
||||
self.ft_token_id = self.vocab_size_with_mem + 2
|
||||
self.ft_token_id = self.vocab_size_with_mem + 2
|
||||
|
||||
self.icae.resize_token_embeddings(self.vocab_size_with_mem + 3)
|
||||
|
||||
self.icae.resize_token_embeddings(self.vocab_size_with_mem + 3)
|
||||
|
||||
# special tokens for Llama-2/Mistral tokenizer
|
||||
self.bos_id = 1
|
||||
self.eos_id = 2
|
||||
|
||||
|
||||
self.dim = self.icae.config.hidden_size
|
||||
self.icae = get_peft_model(self.icae, lora_config)
|
||||
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
self.memory_token_embed = nn.Embedding(self.mem_size + 3, self.dim, padding_idx=None)
|
||||
self.loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_fast=False)
|
||||
self.append_sequence = torch.arange(self.vocab_size, self.vocab_size + self.mem_size, dtype=torch.long, device=device).unsqueeze(0) # mem tokens
|
||||
|
||||
self.append_sequence = torch.arange(
|
||||
self.vocab_size, self.vocab_size + self.mem_size, dtype=torch.long, device=device
|
||||
).unsqueeze(
|
||||
0
|
||||
) # mem tokens
|
||||
|
||||
if self.training:
|
||||
self.init()
|
||||
|
||||
|
||||
def init(self):
|
||||
print("Freezing the decoder...")
|
||||
freeze_model(self.decoder)
|
||||
|
|
@ -151,15 +174,15 @@ class ICAE(torch.nn.Module):
|
|||
print(f"Finished loading from {self.training_args.restore_from}")
|
||||
print("Enabling gradient checkpointing...")
|
||||
# self.icae.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
|
||||
self.decoder.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
|
||||
|
||||
|
||||
self.decoder.gradient_checkpointing_enable(
|
||||
gradient_checkpointing_kwargs={"use_reentrant": False}
|
||||
)
|
||||
|
||||
def compute_num_segments(self, total_length):
|
||||
assert total_length > 0
|
||||
num_segments = math.ceil(total_length / (self.mem_size * self.mean_compression_rate))
|
||||
return num_segments
|
||||
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
|
|
@ -171,73 +194,93 @@ class ICAE(torch.nn.Module):
|
|||
total_length = input_ids.size(1)
|
||||
num_segments = self.compute_num_segments(total_length)
|
||||
segment_length = math.ceil(total_length / num_segments)
|
||||
|
||||
|
||||
prompt_answer_embs = self.icae.get_base_model().model.embed_tokens(prompt_answer_ids)
|
||||
max_compressed_length = num_segments * self.mem_size
|
||||
compress_outputs = torch.zeros((max_compressed_length, self.dim)).to(prompt_answer_embs)
|
||||
|
||||
|
||||
for segment_idx in range(num_segments):
|
||||
|
||||
|
||||
start_idx = segment_idx * segment_length
|
||||
end_idx = min((segment_idx + 1) * segment_length, total_length)
|
||||
segment_input_ids = input_ids[:, start_idx:end_idx]
|
||||
segment_input_ids = torch.cat([segment_input_ids, self.append_sequence], dim=1)
|
||||
mem_flag = segment_input_ids >= self.vocab_size
|
||||
|
||||
segment_input_embedding = self.icae.get_base_model().model.embed_tokens(segment_input_ids)
|
||||
segment_input_embedding[mem_flag] = self.memory_token_embed(segment_input_ids[mem_flag] - self.vocab_size).to(segment_input_embedding)
|
||||
segment_input_embedding = self.icae.get_base_model().model.embed_tokens(
|
||||
segment_input_ids
|
||||
)
|
||||
segment_input_embedding[mem_flag] = self.memory_token_embed(
|
||||
segment_input_ids[mem_flag] - self.vocab_size
|
||||
).to(segment_input_embedding)
|
||||
|
||||
# compress the current segment
|
||||
segment_compress_outputs = self.icae(inputs_embeds=segment_input_embedding, output_hidden_states=True)
|
||||
segment_compress_outputs = self.icae(
|
||||
inputs_embeds=segment_input_embedding, output_hidden_states=True
|
||||
)
|
||||
segment_compress_outputs = segment_compress_outputs.hidden_states[-1]
|
||||
|
||||
# collect memory tokens
|
||||
compress_outputs[segment_idx*self.mem_size: self.mem_size*(segment_idx+1)] = segment_compress_outputs[mem_flag]
|
||||
|
||||
compress_outputs[
|
||||
segment_idx * self.mem_size : self.mem_size * (segment_idx + 1)
|
||||
] = segment_compress_outputs[mem_flag]
|
||||
|
||||
del segment_input_ids, segment_input_embedding
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
# decoder part
|
||||
decoder_mem_flag = (prompt_answer_ids >= self.vocab_size) & (prompt_answer_ids < self.vocab_size + self.mem_size) # only mem tokens
|
||||
decoder_mem_flag = (prompt_answer_ids >= self.vocab_size) & (
|
||||
prompt_answer_ids < self.vocab_size + self.mem_size
|
||||
) # only mem tokens
|
||||
|
||||
prompt_answer_embs[decoder_mem_flag] = compress_outputs # replace memory slots
|
||||
special_prompt = prompt_answer_ids >= self.vocab_size_with_mem
|
||||
prompt_answer_embs[special_prompt] = self.memory_token_embed(prompt_answer_ids[special_prompt] - self.vocab_size).to(prompt_answer_embs) # replace special token's embedding from self.memory_token_embed
|
||||
|
||||
if self.training: # has an independent se.f.decoder
|
||||
decoder_outputs = self.decoder(inputs_embeds=prompt_answer_embs, output_hidden_states=True)
|
||||
else:
|
||||
with self.icae.disable_adapter(): # no independent decoder; use self.icae
|
||||
decoder_outputs = self.icae(inputs_embeds=prompt_answer_embs, output_hidden_states=True)
|
||||
prompt_answer_embs[special_prompt] = self.memory_token_embed(
|
||||
prompt_answer_ids[special_prompt] - self.vocab_size
|
||||
).to(
|
||||
prompt_answer_embs
|
||||
) # replace special token's embedding from self.memory_token_embed
|
||||
|
||||
if self.training: # has an independent se.f.decoder
|
||||
decoder_outputs = self.decoder(
|
||||
inputs_embeds=prompt_answer_embs, output_hidden_states=True
|
||||
)
|
||||
else:
|
||||
with self.icae.disable_adapter(): # no independent decoder; use self.icae
|
||||
decoder_outputs = self.icae(
|
||||
inputs_embeds=prompt_answer_embs, output_hidden_states=True
|
||||
)
|
||||
|
||||
logits = decoder_outputs.logits
|
||||
effective_logits = logits[:,:-1,:].reshape(-1, logits.size(-1))
|
||||
target_ids = labels[:,1:].reshape(-1)
|
||||
effective_logits = logits[:, :-1, :].reshape(-1, logits.size(-1))
|
||||
target_ids = labels[:, 1:].reshape(-1)
|
||||
loss = self.loss_fct(effective_logits, target_ids)
|
||||
return {"loss": loss, "logits": logits}
|
||||
|
||||
|
||||
def tokens_to_embeddings(self, token_ids): # input_tokens can be either normal tokens and special tokens
|
||||
|
||||
def tokens_to_embeddings(
|
||||
self, token_ids
|
||||
): # input_tokens can be either normal tokens and special tokens
|
||||
embeddings = self.icae.get_base_model().model.embed_tokens(token_ids)
|
||||
special_flags = token_ids >= self.vocab_size
|
||||
embeddings[special_flags] = self.memory_token_embed(token_ids[special_flags] - self.vocab_size).to(embeddings) # replace special token's embedding from self.memory_token_embed
|
||||
embeddings[special_flags] = self.memory_token_embed(
|
||||
token_ids[special_flags] - self.vocab_size
|
||||
).to(
|
||||
embeddings
|
||||
) # replace special token's embedding from self.memory_token_embed
|
||||
return embeddings
|
||||
|
||||
|
||||
|
||||
def _compress(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None
|
||||
self, input_ids: torch.LongTensor = None
|
||||
): # for inference; compress a fixed length of input into memory slots
|
||||
|
||||
batch_size = input_ids.size(0)
|
||||
total_length = input_ids.size(1)
|
||||
num_segments = self.compute_num_segments(total_length)
|
||||
segment_length = math.ceil(total_length / num_segments)
|
||||
|
||||
|
||||
max_compressed_length = num_segments * self.mem_size
|
||||
compress_outputs = torch.zeros((max_compressed_length, self.dim))
|
||||
|
||||
|
||||
for segment_idx in range(num_segments):
|
||||
start_idx = segment_idx * segment_length
|
||||
end_idx = min((segment_idx + 1) * segment_length, total_length)
|
||||
|
|
@ -245,17 +288,25 @@ class ICAE(torch.nn.Module):
|
|||
segment_input_ids = torch.cat([segment_input_ids, self.append_sequence], dim=1)
|
||||
mem_flag = segment_input_ids >= self.vocab_size
|
||||
|
||||
segment_input_embedding = self.icae.get_base_model().model.embed_tokens(segment_input_ids)
|
||||
segment_input_embedding[mem_flag] = self.memory_token_embed(segment_input_ids[mem_flag] - self.vocab_size).to(segment_input_embedding)
|
||||
segment_input_embedding = self.icae.get_base_model().model.embed_tokens(
|
||||
segment_input_ids
|
||||
)
|
||||
segment_input_embedding[mem_flag] = self.memory_token_embed(
|
||||
segment_input_ids[mem_flag] - self.vocab_size
|
||||
).to(segment_input_embedding)
|
||||
|
||||
# compress the current segment
|
||||
segment_compress_outputs = self.icae(inputs_embeds=segment_input_embedding, output_hidden_states=True)
|
||||
segment_compress_outputs = self.icae(
|
||||
inputs_embeds=segment_input_embedding, output_hidden_states=True
|
||||
)
|
||||
segment_compress_outputs = segment_compress_outputs.hidden_states[-1]
|
||||
|
||||
# collect memory tokens
|
||||
compress_outputs[segment_idx*self.mem_size: self.mem_size*(segment_idx+1)] = segment_compress_outputs[mem_flag]
|
||||
|
||||
compress_outputs[
|
||||
segment_idx * self.mem_size : self.mem_size * (segment_idx + 1)
|
||||
] = segment_compress_outputs[mem_flag]
|
||||
|
||||
del segment_input_ids, segment_input_embedding
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return compress_outputs
|
||||
|
||||
return compress_outputs
|
||||
|
|
|
|||
|
|
@ -1,10 +1,18 @@
|
|||
import transformers
|
||||
from peft import (
|
||||
LoraConfig,
|
||||
)
|
||||
from datasets import load_dataset
|
||||
from modeling_icae_multi_span import ICAE, ModelArguments, DataArguments, TrainingArguments
|
||||
from training_utils import pretrain_tokenize_function, DataCollatorForDynamicPadding, train_model
|
||||
from modeling_icae_multi_span import (
|
||||
ICAE,
|
||||
DataArguments,
|
||||
ModelArguments,
|
||||
TrainingArguments,
|
||||
)
|
||||
from peft import LoraConfig
|
||||
from training_utils import (
|
||||
DataCollatorForDynamicPadding,
|
||||
pretrain_tokenize_function,
|
||||
train_model,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
|
||||
|
|
@ -12,22 +20,23 @@ def main():
|
|||
|
||||
print(model_args)
|
||||
print(data_args)
|
||||
|
||||
training_args.gradient_checkpointing_kwargs = {"use_reentrant": False} # manually add this argument in the code
|
||||
|
||||
training_args.gradient_checkpointing_kwargs = {
|
||||
"use_reentrant": False
|
||||
} # manually add this argument in the code
|
||||
|
||||
lora_config = LoraConfig(
|
||||
r=model_args.lora_r,
|
||||
lora_alpha=32,
|
||||
lora_dropout=0.05,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM"
|
||||
r=model_args.lora_r, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
|
||||
)
|
||||
|
||||
# check model_args.mem_size and min_tokens_for_lm
|
||||
assert (training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)) == 0, "training_args.fixed_mem_size must be a power of 2"
|
||||
assert training_args.leave_tokens_for_lm <= training_args.min_tokens_for_lm, "leave_tokens_for_lm should be fewer than min_tokens_for_lm"
|
||||
|
||||
|
||||
# check model_args.mem_size and min_tokens_for_lm
|
||||
assert (
|
||||
training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)
|
||||
) == 0, "training_args.fixed_mem_size must be a power of 2"
|
||||
assert (
|
||||
training_args.leave_tokens_for_lm <= training_args.min_tokens_for_lm
|
||||
), "leave_tokens_for_lm should be fewer than min_tokens_for_lm"
|
||||
|
||||
memory_size = training_args.fixed_mem_size
|
||||
|
||||
train_file = "/path/to/train/file"
|
||||
|
|
@ -35,17 +44,27 @@ def main():
|
|||
|
||||
print("Loading dataset...")
|
||||
|
||||
dataset = load_dataset("json", data_files={"train": train_file, "eval": eval_file}, streaming=True) # streaming can be removed if the dataset is not very large.
|
||||
dataset = load_dataset(
|
||||
"json", data_files={"train": train_file, "eval": eval_file}, streaming=True
|
||||
) # streaming can be removed if the dataset is not very large.
|
||||
train_dataset = dataset["train"]
|
||||
eval_dataset = dataset["eval"]
|
||||
|
||||
model = ICAE(model_args, training_args, lora_config)
|
||||
MEM_TOKENS = list(range(model.vocab_size, model.vocab_size + memory_size))
|
||||
|
||||
train_dataset = train_dataset.map(pretrain_tokenize_function, batched=True, batch_size=64, fn_kwargs={"model": model, "mem": MEM_TOKENS, "lm_ratio": training_args.lm_ratio})
|
||||
eval_dataset = eval_dataset.map(pretrain_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS}) # don't add lm in the dev set.
|
||||
train_dataset = train_dataset.map(
|
||||
pretrain_tokenize_function,
|
||||
batched=True,
|
||||
batch_size=64,
|
||||
fn_kwargs={"model": model, "mem": MEM_TOKENS, "lm_ratio": training_args.lm_ratio},
|
||||
)
|
||||
eval_dataset = eval_dataset.map(
|
||||
pretrain_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS}
|
||||
) # don't add lm in the dev set.
|
||||
|
||||
data_collator = DataCollatorForDynamicPadding(model.pad_token_id)
|
||||
train_model(model, train_dataset, eval_dataset, training_args, data_collator)
|
||||
|
||||
main()
|
||||
|
||||
main()
|
||||
|
|
|
|||
|
|
@ -1,11 +1,17 @@
|
|||
import json
|
||||
import sys
|
||||
|
||||
import torch
|
||||
from modeling_icae_multi_span import (
|
||||
ICAE,
|
||||
DataArguments,
|
||||
ModelArguments,
|
||||
TrainingArguments,
|
||||
)
|
||||
from peft import LoraConfig
|
||||
from safetensors.torch import load_file
|
||||
from tqdm import tqdm
|
||||
from transformers import HfArgumentParser
|
||||
from peft import LoraConfig
|
||||
from modeling_icae_multi_span import ICAE, ModelArguments, DataArguments, TrainingArguments
|
||||
import sys
|
||||
from safetensors.torch import load_file
|
||||
|
||||
# Set the computation device
|
||||
device = "cuda"
|
||||
|
|
@ -16,11 +22,7 @@ model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
|||
|
||||
# Define Lora configuration
|
||||
lora_config = LoraConfig(
|
||||
r=512,
|
||||
lora_alpha=32,
|
||||
lora_dropout=model_args.lora_dropout,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM"
|
||||
r=512, lora_alpha=32, lora_dropout=model_args.lora_dropout, bias="none", task_type="CAUSAL_LM"
|
||||
)
|
||||
|
||||
# Initialize model and send it to CUDA device
|
||||
|
|
@ -29,11 +31,13 @@ model = ICAE(model_args, training_args, lora_config)
|
|||
# Load the fine-tuned checkpoint
|
||||
print(f"Loading trained checkpoint from {training_args.output_dir}")
|
||||
state_dict = load_file(training_args.output_dir)
|
||||
model.load_state_dict(state_dict, strict=False) # only load lora and memory token embeddings
|
||||
model.load_state_dict(state_dict, strict=False) # only load lora and memory token embeddings
|
||||
|
||||
model = model.to(device)
|
||||
|
||||
lines = ["I don't have a favorite condiment as I don't consume food or condiments. However, I can tell you that many people enjoy condiments like ketchup, mayonnaise, mustard, soy sauce, hot sauce, and ranch dressing, among others. The favorite condiment can vary greatly from person to person, depending on their taste preferences and cultural influences."]
|
||||
lines = [
|
||||
"I don't have a favorite condiment as I don't consume food or condiments. However, I can tell you that many people enjoy condiments like ketchup, mayonnaise, mustard, soy sauce, hot sauce, and ranch dressing, among others. The favorite condiment can vary greatly from person to person, depending on their taste preferences and cultural influences."
|
||||
]
|
||||
|
||||
# Prepare the model for evaluation
|
||||
model.eval()
|
||||
|
|
@ -42,21 +46,23 @@ with torch.no_grad():
|
|||
|
||||
for line in tqdm(lines):
|
||||
# Tokenize input text
|
||||
tokenized_text = model.tokenizer(line, truncation=True,
|
||||
max_length=5120, padding=False,
|
||||
return_attention_mask=False)
|
||||
tokenized_text = model.tokenizer(
|
||||
line, truncation=True, max_length=5120, padding=False, return_attention_mask=False
|
||||
)
|
||||
# Generate compressed outputs
|
||||
input_ids = torch.LongTensor([tokenized_text['input_ids']]).to(device)
|
||||
input_ids = torch.LongTensor([tokenized_text["input_ids"]]).to(device)
|
||||
memory_slots = model._compress(input_ids)
|
||||
|
||||
|
||||
# prompt_output = model.tokenizer(data['prompt'], add_special_tokens=False, padding=False)
|
||||
prompt_ids = torch.LongTensor([[model.ae_token_id]]).to(device)
|
||||
|
||||
prompt_answer_embs = model.tokens_to_embeddings(prompt_ids)
|
||||
memory_slots = memory_slots.to(prompt_answer_embs)
|
||||
|
||||
|
||||
# Concatenate and clone input embeddings
|
||||
decoder_input_embeddings = torch.cat((memory_slots.unsqueeze(0), prompt_answer_embs), dim=1)
|
||||
decoder_input_embeddings = torch.cat(
|
||||
(memory_slots.unsqueeze(0), prompt_answer_embs), dim=1
|
||||
)
|
||||
output = decoder_input_embeddings.clone()
|
||||
|
||||
generate_text = []
|
||||
|
|
@ -64,25 +70,30 @@ with torch.no_grad():
|
|||
|
||||
# Generate text output
|
||||
for i in range(512):
|
||||
with model.icae.disable_adapter(): # no independent decoder; use self.icae
|
||||
out = model.icae(inputs_embeds=output, past_key_values=past_key_values, use_cache=True)
|
||||
logit = out.logits[:, -1, :model.vocab_size-1]
|
||||
with model.icae.disable_adapter(): # no independent decoder; use self.icae
|
||||
out = model.icae(
|
||||
inputs_embeds=output, past_key_values=past_key_values, use_cache=True
|
||||
)
|
||||
logit = out.logits[:, -1, : model.vocab_size - 1]
|
||||
past_key_values = out.past_key_values
|
||||
|
||||
next_token_id = torch.argmax(logit, dim=-1)
|
||||
# print(next_token_id)
|
||||
|
||||
if next_token_id.item() == 2: # eos
|
||||
|
||||
if next_token_id.item() == 2: # eos
|
||||
break
|
||||
|
||||
output = model.icae.get_base_model().model.embed_tokens(next_token_id).unsqueeze(1).to(device)
|
||||
output = (
|
||||
model.icae.get_base_model()
|
||||
.model.embed_tokens(next_token_id)
|
||||
.unsqueeze(1)
|
||||
.to(device)
|
||||
)
|
||||
generate_text.append(next_token_id.item())
|
||||
|
||||
generated_text = model.tokenizer.decode(generate_text)
|
||||
|
||||
# Structure output data
|
||||
output_ = {
|
||||
"output": generated_text
|
||||
}
|
||||
output_ = {"output": generated_text}
|
||||
|
||||
f.write(json.dumps(output_) + "\n")
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
from transformers import Trainer
|
||||
import math
|
||||
import os
|
||||
import torch
|
||||
import random
|
||||
|
||||
import torch
|
||||
from transformers import Trainer
|
||||
from transformers.trainer_utils import get_last_checkpoint
|
||||
import math
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue