commit 404566501bbb6ab539679cbc2ff1e6c2ed257a8e Author: 51616 Date: Wed Dec 18 13:05:19 2024 +0000 init commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4905ebf --- /dev/null +++ b/.gitignore @@ -0,0 +1,14 @@ +tmp/ +models/ +results/ +datasets/ +data/ +__pycache__/ +*egg-info/ +.vscode/ +.ipynb_checkpoints/ +wandb/ +data/ +*.bak +*outputs/ +plots/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..dd27b51 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,17 @@ +repos: + - repo: https://github.com/asottile/pyupgrade + rev: v2.31.1 + hooks: + - id: pyupgrade + args: + - --py39-plus + - repo: https://github.com/PyCQA/isort + rev: 5.12.0 + hooks: + - id: isort + args: + - --profile=black + - repo: https://github.com/python/black + rev: 22.3.0 + hooks: + - id: black diff --git a/hyperlora/instruction_finetune.py b/hyperlora/instruction_finetune.py new file mode 100644 index 0000000..6cd072a --- /dev/null +++ b/hyperlora/instruction_finetune.py @@ -0,0 +1,49 @@ +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 + +def main(): + parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) + model_args, data_args, training_args = parser.parse_args_into_dataclasses() + + print(model_args) + print(data_args) + + 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" + ) + + # 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" + + memory_size = training_args.fixed_mem_size + + train_file = "/path/to/train/file" + eval_file = "/path/to/dev/file" + + print("Loading dataset...") + + dataset = load_dataset("json", data_files={"train": train_file, "eval": eval_file}) + + 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(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() \ No newline at end of file diff --git a/icae_v2/__init__.py b/icae_v2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/icae_v2/dev_v2.jsonl b/icae_v2/dev_v2.jsonl new file mode 100644 index 0000000..e071b78 --- /dev/null +++ b/icae_v2/dev_v2.jsonl @@ -0,0 +1,6 @@ +{"input": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "Identify the person arrested on suspicion of spying for North Korea.", "answer": "Benoit Quennedey"} +{"input": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "List the actions taken against Benoit Quennedey after his arrest.", "answer": "His office in the French Senate was raided by DGSI officers."} +{"input": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "What is the role of Benoit Quennedey in the French government?", "answer": "He is a senior civil servant who liaises between the French Senate and the Department of Architecture and Heritage."} +{"input": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "What are the charges against Benoit Quennedey?", "answer": "He is suspected of collecting and delivering to a foreign power information likely to subvert core national interests."} +{"input": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "Mention the organization Benoit Quennedey is believed to be the president of.", "answer": "Franco-Korean Friendship Association"} +{"input": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "When did the counterintelligence investigation into Quennedey's activities begin?", "answer": "In March of this year."} \ No newline at end of file diff --git a/icae_v2/fine_tuned_inference.py b/icae_v2/fine_tuned_inference.py new file mode 100644 index 0000000..c675475 --- /dev/null +++ b/icae_v2/fine_tuned_inference.py @@ -0,0 +1,102 @@ +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 +from safetensors.torch import load_file + +# Set the computation device +device = "cuda" + +# Parse model, data, and training arguments +parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) +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" +) + +# Initialize model and send it to CUDA device +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 = model.to(device) + +# Read the data file +file_path = "./dev_v2.jsonl" +lines = None +with open(file_path, "r") as f: + lines = f.readlines() + +# Prepare the model for evaluation +max_out_length = 512 +model.eval() + +with torch.no_grad(): + with open("ft_inference.out", "w") as f: + + 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) + # Generate compressed outputs + 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_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) + output = decoder_input_embeddings.clone() + + generate_text = [] + past_key_values = None + + # 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) + # out = decoder(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 + break + + 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'], + "prompt": data["prompt"], + "output": generated_text, + "answer": data["answer"] + } + + f.write(json.dumps(output_) + "\n") diff --git a/icae_v2/fine_tuned_inference_script.sh b/icae_v2/fine_tuned_inference_script.sh new file mode 100644 index 0000000..9de96f7 --- /dev/null +++ b/icae_v2/fine_tuned_inference_script.sh @@ -0,0 +1,16 @@ +#!/bin/bash + +# MODEL="mistralai/Mistral-7B-v0.1" +BASE_MODEL="mistralai/Mistral-7B-Instruct-v0.2" +# MODEL="meta-llama/Llama-2-7b-hf" +# MODEL="meta-llama/Llama-2-7b-chat-hf" +MODEL_NAME="${MODEL//\//-}" + +maxlen=5120 +mem=128 +r=512 +mean_compression_rate=4 + +ICAE_MODEL_PATH=$1 # ICAE model to use; wget "https://huggingface.co/sggetao/icae/resolve/main/mistral_7b_ft_icae.safetensors" + +python fine_tuned_inference.py --mean_compression_rate $mean_compression_rate --model_max_length $maxlen --fixed_mem_size $mem --lora_r $r --output_dir $ICAE_MODEL_PATH --model_name_or_path $BASE_MODEL --bf16 --train False diff --git a/icae_v2/instruction_finetune.py b/icae_v2/instruction_finetune.py new file mode 100644 index 0000000..6cd072a --- /dev/null +++ b/icae_v2/instruction_finetune.py @@ -0,0 +1,49 @@ +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 + +def main(): + parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) + model_args, data_args, training_args = parser.parse_args_into_dataclasses() + + print(model_args) + print(data_args) + + 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" + ) + + # 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" + + memory_size = training_args.fixed_mem_size + + train_file = "/path/to/train/file" + eval_file = "/path/to/dev/file" + + print("Loading dataset...") + + dataset = load_dataset("json", data_files={"train": train_file, "eval": eval_file}) + + 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(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() \ No newline at end of file diff --git a/icae_v2/modeling_icae_multi_span.py b/icae_v2/modeling_icae_multi_span.py new file mode 100644 index 0000000..3da71d8 --- /dev/null +++ b/icae_v2/modeling_icae_multi_span.py @@ -0,0 +1,261 @@ +# ICAE that supports multi span concat + +import transformers +from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig +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 +from safetensors.torch import load_file + +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"} + ) + train: bool = field( + default=True, + 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"}) + +@dataclass +class TrainingArguments(transformers.TrainingArguments): + cache_dir: Optional[str] = field(default=None) + 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)."}, + ) + fixed_mem_size: int = field( + default=128, + metadata={"help": "Enalbing the fixed mem size."}, + ) + mean_compression_rate: int = field( + default=4, + metadata={"help": "Mean compression rate; default=4"}, + ) + min_tokens_for_lm: int = field( + default=64, + metadata={"help": "Minimum tokens for lm objective learning"}, + ) + leave_tokens_for_lm: int = field( + default=8, + metadata={"help": "Leave some tokens without loss for lm objective"}, + ) + lm_ratio: float = field( + default=0.0, + metadata={"help": "Ratio for LM training."}, + ) + add_special_token_for_lm: bool = field( + 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"} + ) + +def print_trainable_parameters(model): + trainable_parameters = 0 + all_param = 0 + for _, param in model.named_parameters(): + 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}") + # for name, param in model.named_parameters(): + # if param.requires_grad: + # print(name, param.shape) + + +def freeze_model(model): + for _, param in model.named_parameters(): + param.requires_grad = False + + +class ICAE(torch.nn.Module): + def __init__(self, model_args, training_args, lora_config): + super().__init__() + 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.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) + + # 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.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 + + if self.training: + self.init() + + + def init(self): + print("Freezing the decoder...") + freeze_model(self.decoder) + self.decoder.eval() + print_trainable_parameters(self) + if self.training_args.restore_from is not None and self.training_args.restore_from != "": + print(f"Loading from the pretrained checkpoint: {self.training_args.restore_from}...") + state_dict = load_file(self.training_args.restore_from) + self.load_state_dict(state_dict) + 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}) + + + 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, + prompt_answer_ids: torch.LongTensor = None, + labels: Optional[torch.LongTensor] = None, + ): + # encoder part + 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) + + 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) + + # compress the current segment + 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] + + 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 + + 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) + + + logits = decoder_outputs.logits + 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 + 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 + return embeddings + + + def _compress( + 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) + 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) + + # compress the current segment + 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] + + del segment_input_ids, segment_input_embedding + torch.cuda.empty_cache() + + return compress_outputs \ No newline at end of file diff --git a/icae_v2/pretrain.py b/icae_v2/pretrain.py new file mode 100644 index 0000000..ed78a7a --- /dev/null +++ b/icae_v2/pretrain.py @@ -0,0 +1,51 @@ +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 + +def main(): + parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) + model_args, data_args, training_args = parser.parse_args_into_dataclasses() + + print(model_args) + print(data_args) + + 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" + ) + + # 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" + eval_file = "/path/to/dev/file" + + 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. + 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. + + data_collator = DataCollatorForDynamicPadding(model.pad_token_id) + train_model(model, train_dataset, eval_dataset, training_args, data_collator) + +main() \ No newline at end of file diff --git a/icae_v2/pretrained_inference.py b/icae_v2/pretrained_inference.py new file mode 100644 index 0000000..68271a2 --- /dev/null +++ b/icae_v2/pretrained_inference.py @@ -0,0 +1,88 @@ +import json +import torch +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" + +# Parse model, data, and training arguments +parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) +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" +) + +# Initialize model and send it to CUDA device +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 = 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."] + +# Prepare the model for evaluation +model.eval() +with torch.no_grad(): + with open("tmp.out", "w") as f: + + for line in tqdm(lines): + # Tokenize input text + 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) + 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) + output = decoder_input_embeddings.clone() + + generate_text = [] + past_key_values = None + + # 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] + 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 + break + + 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 + } + + f.write(json.dumps(output_) + "\n") diff --git a/icae_v2/pretrained_inference_script.sh b/icae_v2/pretrained_inference_script.sh new file mode 100644 index 0000000..f8d7883 --- /dev/null +++ b/icae_v2/pretrained_inference_script.sh @@ -0,0 +1,16 @@ +#!/bin/bash + +# MODEL="mistralai/Mistral-7B-v0.1" +BASE_MODEL="mistralai/Mistral-7B-Instruct-v0.2" +# MODEL="meta-llama/Llama-2-7b-hf" +# MODEL="meta-llama/Llama-2-7b-chat-hf" +MODEL_NAME="${MODEL//\//-}" + +maxlen=5120 +mem=128 +r=512 +mean_compression_rate=4 + +ICAE_MODEL_PATH=$1 # ICAE model to use; wget "https://huggingface.co/sggetao/icae/resolve/main/mistral_7b_pretrained_icae.safetensors" + +python pretrained_inference.py --mean_compression_rate $mean_compression_rate --model_max_length $maxlen --fixed_mem_size $mem --lora_r $r --output_dir $ICAE_MODEL_PATH --model_name_or_path $BASE_MODEL --bf16 --train False diff --git a/icae_v2/training_utils.py b/icae_v2/training_utils.py new file mode 100644 index 0000000..bfa5d94 --- /dev/null +++ b/icae_v2/training_utils.py @@ -0,0 +1,219 @@ +from transformers import Trainer +import os +import torch +import random + +from transformers.trainer_utils import get_last_checkpoint +import math + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def train_model(model, train_dataset, eval_dataset, training_args, data_collator=None): + + last_checkpoint = None + if os.path.isdir(training_args.output_dir) and not training_args.overwrite_output_dir: + last_checkpoint = get_last_checkpoint(training_args.output_dir) + if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0: + raise ValueError( + f"Output directory ({training_args.output_dir}) already exists and is not empty. " + "Use --overwrite_output_dir to overcome." + ) + elif last_checkpoint is not None and training_args.resume_from_checkpoint is None: + print( + f"Checkpoint detected, resuming training at {last_checkpoint}. " + "To avoid this behavior, change " + "the `--output_dir` or add `--overwrite_output_dir` to train from scratch." + ) + + if ( + max(training_args.per_device_train_batch_size, training_args.per_device_eval_batch_size) + == 1 + ): + data_collator = None + + # print training_args at local_rank 0 + local_rank = int(os.getenv("LOCAL_RANK", "0")) + if local_rank == 0: + print(training_args) + + trainer = Trainer( + model=model, + args=training_args, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + data_collator=data_collator, + ) + + checkpoint = None + + if training_args.resume_from_checkpoint is not None: + checkpoint = training_args.resume_from_checkpoint + elif last_checkpoint is not None: + checkpoint = last_checkpoint + + print(f"Loaded from the checkpoint: {checkpoint}") + + train_result = trainer.train(resume_from_checkpoint=checkpoint) + trainer.save_model() + trainer.log_metrics("train", train_result.metrics) + metrics = trainer.evaluate() + trainer.log_metrics("eval", metrics) + trainer.save_metrics("eval", metrics) + + +def text_extraction(input_ids, length, lm_ratio=0.0): + + input_len = len(input_ids) + assert input_len >= 1, f"Error: invalid input length ({input_len})" + + # ae + if random.random() >= lm_ratio: + if input_len <= length: # if shorter, keep the complete text + return input_ids, [] + else: + last_start = input_len - length + random_start = random.randint(0, last_start) + return input_ids[random_start : random_start + length], [] + + # lm + if input_len <= length: + r = random.randint(0, input_len - 1) + return input_ids[: r + 1], input_ids[r + 1 :] + else: + last_start = input_len - length + random_start = random.randint(0, last_start) + return input_ids[random_start : random_start + length], input_ids[random_start + length :] + + +def pretrain_tokenize_function(examples, model, mem, lm_ratio=0.0): + text_output = model.tokenizer( + examples["text"], truncation=False, padding=False, return_attention_mask=False + ) + text_output["prompt_answer_ids"] = [] + text_output["labels"] = [] + + max_len = model.training_args.model_max_length # heuristic + + for idx in range(len(text_output["input_ids"])): + + ae = True + a, b = text_extraction(text_output["input_ids"][idx], max_len, lm_ratio=lm_ratio) + length_a = len(a) + num_segments = model.compute_num_segments(length_a) + total_mem_length = num_segments * model.mem_size + + if ( + len(b) > model.training_args.min_tokens_for_lm + ): # avoid too few tokens for lm, which is a waste of computing + ae = False + b = b[:max_len] + + text_output["input_ids"][idx] = a + + # decoder part: note that in v2, we add mem_tokens to the prompt_ids + # for easy implementation; which is different from v1 implementation + # where mem tokens are not in the prompt_ids + if ae: # autoencoding objective + prompt_ids = [mem[0]] * total_mem_length + [model.ae_token_id] + answer_ids = a + [model.eos_id] # if ae, eos token + else: # lm objective + prompt_ids = [mem[0]] * total_mem_length + if model.training_args.add_special_token_for_lm: + prompt_ids += [model.lm_token_id] + answer_ids = b # if lm, no eos token + + text_output["prompt_answer_ids"].append(prompt_ids + answer_ids) + if ae: + labels = [-100] * len(prompt_ids) + answer_ids + else: + labels = ( + [-100] * len(prompt_ids) + + [-100] * model.training_args.leave_tokens_for_lm + + answer_ids[model.training_args.leave_tokens_for_lm :] + ) # no loss for leave_tokens_for_lm + text_output["labels"].append(labels) + assert len(text_output["prompt_answer_ids"][-1]) == len(labels) + + return text_output + + +def instruct_ft_tokenize_function(examples, model, mem): + text_output = model.tokenizer( + examples["input"], + max_length=5120, + truncation=True, + padding=False, + return_attention_mask=False, + add_special_tokens=False, + ) + prompt_output = model.tokenizer( + examples["prompt"], + truncation=False, + padding=False, + return_attention_mask=False, + add_special_tokens=False, + ) + label_output = model.tokenizer( + examples["answer"], + truncation=False, + padding=False, + return_attention_mask=False, + add_special_tokens=False, + ) + text_output["prompt_answer_ids"] = [] + text_output["labels"] = [] + + max_len = model.training_args.model_max_length # heuristic + + for idx in range(len(text_output["input_ids"])): + + length = len(text_output["input_ids"][idx]) + num_segments = model.compute_num_segments(length) + total_mem_length = num_segments * model.mem_size + + prompt_ids = ( + [mem[0]] * total_mem_length + [model.ft_token_id] + prompt_output["input_ids"][idx] + ) + prompt_ids = ( + [1, 733, 16289, 28793] + prompt_ids + [733, 28748, 16289, 28793] + ) # special formats for prompt in Mistral + answer_ids = label_output["input_ids"][idx] + [model.eos_id] + + text_output["prompt_answer_ids"].append(prompt_ids + answer_ids) + + labels = [-100] * len(prompt_ids) + answer_ids + text_output["labels"].append(labels) + + assert len(text_output["prompt_answer_ids"][-1]) == len(labels) + + return text_output + + +class DataCollatorForDynamicPadding: + def __init__(self, pad_token_id, pad_to_multiple_of=None): + self.pad_token_id = pad_token_id + self.pad_to_multiple_of = pad_to_multiple_of + + def __call__(self, examples): + input_ids = [torch.tensor(example["input_ids"], dtype=torch.long) for example in examples] + labels = [torch.tensor(example["labels"], dtype=torch.long) for example in examples] + prompt_answer_ids = [ + torch.tensor(example["prompt_answer_ids"], dtype=torch.long) for example in examples + ] + input_ids = self.dynamic_padding(input_ids, fill_value=self.pad_token_id) + prompt_answer_ids = self.dynamic_padding(prompt_answer_ids, fill_value=self.pad_token_id) + labels = self.dynamic_padding(labels) + batch = {"input_ids": input_ids, "labels": labels, "prompt_answer_ids": prompt_answer_ids} + return batch + + def dynamic_padding(self, sequences, fill_value=-100): + max_length = max(len(x) for x in sequences) + if self.pad_to_multiple_of: + max_length = ( + (max_length - 1) // self.pad_to_multiple_of + 1 + ) * self.pad_to_multiple_of + padded_sequences = torch.full((len(sequences), max_length), fill_value, dtype=torch.long) + for i, seq in enumerate(sequences): + padded_sequences[i, : len(seq)] = seq + return padded_sequences diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..4af129a --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,39 @@ +[project] +name = "ctx-to-lora" +version = "0.0.1" +authors = [{ name = "Rujikorn Charakorn" }] +description = "" +readme = "README.md" +requires-python = ">= 3.10" +dependencies = [ + "transformers", + "peft", + "accelerate", + "datasets", + "bitsandbytes", + "jupyter", + "matplotlib", + "hf_transfer", + "torchmetrics", + "inflect", + "pre-commit", + "trl", + "tensorboardX", + "wandb", + "fasttext-wheel", +] + +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[tool.isort] +profile = "black" + + +[tool.black] +line-length = 99 + +[tool.pyright] +exclude = ["train_outputs", "data", "generated_tasks", "outputs", "plots", "tmp", "wandb"] +typeCheckingMode = "off"