diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py
new file mode 100644
index 0000000..a99fd24
--- /dev/null
+++ b/hyperlora/intx_sft.py
@@ -0,0 +1,135 @@
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    HfArgumentParser,
+    TrainingArguments,
+    DataCollatorForSeq2Seq,
+)
+
+from modeling_utils import ModulatedPretrainedModel
+from training_utils import train_model
+
+
+def compute_metrics(eval_pred) -> dict:
+    """
+    Custom metrics function for the trainer
+    Args:
+        eval_pred: tuple of predictions and labels
+    Returns:
+        dictionary containing metric names (str) and values (Any)
+    """
+
+    # preds, labels = eval_preds
+
+    # # predictions is generated tokens for Seq2SeqTrainer
+    # # decode preds and labels
+    # labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+    # decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
+    # decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+    # compute per token accuracy
+    predictions, labels = eval_pred.predictions, eval_pred.label_ids
+
+    # predictions is logits for Trainer
+    preds = predictions.argmax(-1)
+    acc = (preds == labels).mean()
+    return {"per_token_acc": acc}
+
+
+def main():
+    parser = HfArgumentParser((TrainingArguments,))
+    training_args, *_ = parser.parse_args_into_dataclasses()
+
+    training_args.eval_on_start = True
+    training_args.eval_strategy = "steps"
+    training_args.eval_steps = 500
+    training_args.save_strategy = "no"
+    # training_args.save_steps = 500
+    training_args.logging_strategy = "steps"
+    training_args.logging_steps = 100
+
+    # seq2seq args for generation evaluation
+    # training_args.predict_with_generate = True
+    # training_args.generation_max_length = 100
+
+    training_args.gradient_checkpointing_kwargs = {
+        "use_reentrant": False
+    }  # manually add this argument in the code
+
+    model = AutoModelForCausalLM.from_pretrained(
+        "meta-llama/Llama-3.1-8B-Instruct",
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+    )
+    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+    tokenizer.padding_side = "right"
+
+    if isinstance(model, ModulatedPretrainedModel):
+
+        def tokenize(example):
+            model_inputs = tokenizer(
+                example["prompt"],
+                truncation=True,
+                padding=False,
+            )
+            model_inputs["ctx_ids"] = tokenizer(example["context"]).input_ids
+            model_inputs["ctx_attention_mask"] = tokenizer(example["context"]).attention_mask
+            model_inputs["labels"] = ...
+            return model_inputs
+
+    else:
+
+        def tokenize(example):
+            inp = [
+                ctx + "\n" + prompt for ctx, prompt in zip(example["context"], example["prompt"])
+            ]
+            model_inputs = tokenizer(
+                inp,
+                example["answer"],
+                # add_special_tokens=True, ???
+                truncation=True,
+                padding=False,
+            )
+
+            input_ids = model_inputs["input_ids"]
+            labels = [None] * len(input_ids)
+
+            for i in range(len(input_ids)):
+                sequence_ids = model_inputs.sequence_ids(i)
+                labels[i] = [
+                    -100 if sequence_id == 0 else label
+                    for sequence_id, label in zip(sequence_ids, input_ids[i])
+                ]
+            model_inputs["labels"] = labels
+            return model_inputs
+
+    print("Loading dataset...")
+    train_file = "../data/raw_datasets/context_numbers/train.jsonl"
+    eval_file = "../data/raw_datasets/context_numbers/val.jsonl"
+    dataset = load_dataset("json", data_files={"train": train_file, "eval": eval_file})
+    train_ds = dataset["train"].map(tokenize, batched=True)
+    eval_ds = {
+        "train": dataset["train"].select(range(100)).map(tokenize, batched=True),
+        "val": dataset["eval"].map(tokenize, batched=True),
+    }
+
+    # DataCollatorForSeq2Seq also pads the `labels`
+    # useful when we're computing the labels manually
+    # or masking the loss only on completion
+    data_collator = DataCollatorForSeq2Seq(tokenizer, model=model, pad_to_multiple_of=8)
+
+    train_model(
+        model,
+        train_ds,
+        eval_ds,
+        training_args,
+        data_collator,
+        compute_metrics,
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py
new file mode 100644
index 0000000..dd56345
--- /dev/null
+++ b/hyperlora/modeling_utils.py
@@ -0,0 +1,82 @@
+from dataclasses import dataclass, field
+from typing import Optional, Union, Tuple, Any
+
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+# @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"
+#         },
+#     )
+
+
+class ModulatedPretrainedModel(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.lora_config = lora_config
+    #     self.model = 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,
+    #     )
+    def __init__(self, base_model: PreTrainedModel):
+        super().__init__()
+        self.base_model = base_model
+        # NOTE: we can reduce extractor_lm's depth by
+        # e.g., self.extractor_lm.model.layers = self.extractor_lm.model.layers[:num_extractor_layers]
+        # or just use the token embeddings extractor_lm.embed_tokens(input_ids)
+        self.extractor_lm = ...
+        self.hypernet = ...
+
+    def forward(
+        self,
+        ctx_ids: Optional[torch.LongTensor] = None,
+        ctx_attention_mask: Optional[torch.LongTensor] = None,
+        **model_inputs_kwargs: dict[str, Any],
+    ) -> Union[Tuple, ModelOutput]:
+        """Forward pass of the modulated model.
+
+        Args:
+            ctx_ids (torch.LongTensor, optional): Token IDs to be input to the hypernet
+                for generating LoRA parameters. Shape: (batch_size, ctx_length)
+            input_ids (torch.LongTensor, optional): Token IDs to be input to the base
+                language model. Shape: (batch_size, sequence_length)
+            labels (torch.LongTensor, optional): Labels for computing the language modeling
+                loss. Shape: (batch_size, sequence_length)
+
+        Returns:
+            dict: Dictionary containing:
+                - loss (torch.Tensor): Language modeling loss if labels are provided
+                - logits (torch.Tensor): Output logits from the model
+        """
+        if ctx_ids is None:
+            model_outputs = self.base_model(**model_inputs_kwargs)
+            # model_outputs.generated_loras = None
+            return model_outputs
+
+        loss = ...
+        logits = ...
+        generated_loras = ...
+
+        # apply lora hook to the base model
+        self.apply_lora_hook(generated_loras)
+        model_outputs = self.base_model(**model_inputs_kwargs)
+        # model_outputs.generated_loras = generated_loras
+
+        return model_outputs
+
+    def apply_lora_hook(self, generated_loras):
+        pass
diff --git a/hyperlora/training_utils.py b/hyperlora/training_utils.py
new file mode 100644
index 0000000..c76b628
--- /dev/null
+++ b/hyperlora/training_utils.py
@@ -0,0 +1,243 @@
+import math
+import os
+import random
+
+import torch
+from transformers import Seq2SeqTrainer, Trainer
+from transformers.trainer_utils import get_last_checkpoint
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def train_model(
+    model,
+    train_dataset,
+    eval_dataset,
+    training_args,
+    data_collator=None,
+    compute_metrics=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)
+
+    # Seq2SeqTrainer is actually just the same as Trainer
+    # (although it uses a different data collator, i.e., explicit prompt/answer separation)
+    # it just allows `predict_with_generate`
+    # allowing us to compute metrics on the generated outputs
+    # no clue why they call this seq2seq...
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+    )
+
+    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["context"],
+        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/icae_v2/dev_v2.jsonl b/icae_v2/dev_v2.jsonl
index e071b78..8ffbfe6 100644
--- a/icae_v2/dev_v2.jsonl
+++ b/icae_v2/dev_v2.jsonl
@@ -1,6 +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
+{"context": "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"}
+{"context": "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."}
+{"context": "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."}
+{"context": "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."}
+{"context": "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"}
+{"context": "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/training_utils.py b/icae_v2/training_utils.py
index 8227b8e..ac7c0ae 100644
--- a/icae_v2/training_utils.py
+++ b/icae_v2/training_utils.py
@@ -165,7 +165,7 @@ def pretrain_tokenize_function(examples, model, mem, lm_ratio=0.0):
 
 def instruct_ft_tokenize_function(examples, model, mem):
     text_output = model.tokenizer(
-        examples["input"],
+        examples["context"],
         max_length=5120,
         truncation=True,
         padding=False,