{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceType":"datasetVersion","sourceId":15356680,"datasetId":9822808,"databundleVersionId":16266995}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Note\n\nWe explicitely point out which part of the code was taken or was largely inspired by another notebook. If no disclaimer is present, the code was written by us.","metadata":{}},{"cell_type":"code","source":"# *** OUR CODE STARTS HERE ***\n# Importing libraries\nimport pandas as pd\nimport torch\nfrom datasets import Dataset, concatenate_datasets\nfrom transformers import (\n    AutoTokenizer,\n    AutoModelForSeq2SeqLM,\n    DataCollatorForSeq2Seq,\n    Seq2SeqTrainer,\n    Seq2SeqTrainingArguments,\n    TrainerCallback\n)\nimport shutil\nfrom torch.utils.data import DataLoader, IterableDataset\nimport numpy as np\nfrom itertools import cycle\nfrom dataclasses import dataclass\n# *** OUR CODE ENDS HERE ***","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# *** OUR CODE STARTS HERE ***\n# Device checking\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"device: {device}\")\n\nif device.type == \"cuda\":\n    print(f\"name: {torch.cuda.get_device_name(0)}\")\n    print(f\"memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB\")\n# *** OUR CODE ENDS HERE ***","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# *** OUR CODE STARTS HERE ***\n# Loading the model\n# In this case, we fine-tune on ByT5-base continually pre-trained for 2 epochs\n# Requires \"Internet\" on\n\nMODEL_CHECKPOINT = \"jaycentg/akkadian-new-base-cpt-unnormalized\"\ntokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT)\nmodel = AutoModelForSeq2SeqLM.from_pretrained(MODEL_CHECKPOINT, device_map='auto')\n# *** OUR CODE ENDS HERE ***","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# *** OUR CODE STARTS HERE ***\n# Preparing the datasets (challenge_data + top N of extra_data + auxialiary datasets)\n\ndf_add = pd.read_csv(\"/kaggle/input/datasets/marcellogalisai/processed-data/extra_data_processed_ordered.csv\")\n\n# In this case, we use top 7,500 samples\nNUM_SAMPLES = 7500\n\ndf_add = df_add.iloc[:NUM_SAMPLES, :]\n\nprint(df_add.head())\n\ndf_ori = pd.read_csv(\"/kaggle/input/datasets/marcellogalisai/processed-data/challenge_data_processed.csv\")\ndf_ori = df_ori.drop(columns = [\"oare_id\"])\nprint(df_ori.head())\n\ndf_multitask = pd.read_csv(\"/kaggle/input/datasets/marcellogalisai/processed-data/mtl_data_ordered.csv\")\n\nMAX_LENGTH = 512\n\ndef prepare_all_data(df_ori, df_add, df_multitask):\n    val_df   = df_ori.sample(frac=0.1, random_state=42)\n    train_df = df_ori.drop(val_df.index)\n\n    def to_input_target(df, input_col, target_col, prefix):\n        return pd.DataFrame({\n            'input':  prefix + ': ' + df[input_col].astype(str),\n            'target': df[target_col].astype(str),\n        })\n\n    # Setup for main task\n    df_nmt_train = to_input_target(train_df, 'transliteration', 'translation', 'translate Akkadian to English')\n    df_nmt_val = to_input_target(val_df,   'transliteration', 'translation', 'translate Akkadian to English')\n    df_nmt_add = to_input_target(df_add,   'transliteration', 'translation', 'translate Akkadian to English')\n    \n    df_multitask = df_multitask[:(len(df_nmt_train) + len(df_nmt_add))*3]\n\n    # Additional tasks: POS Tagging, NER, Lemmatization\n    df_shuffled = df_multitask.sample(frac=1, random_state=42).reset_index(drop=True)\n    split_size = len(df_shuffled) // 3\n    df_pos_raw = df_shuffled.iloc[:split_size].reset_index(drop=True)\n    df_ner_raw = df_shuffled.iloc[split_size:split_size*2].reset_index(drop=True)\n    df_lemma_raw = df_shuffled.iloc[split_size*2:].reset_index(drop=True)\n\n    df_pos   = to_input_target(df_pos_raw,   'phrase', 'pos',           'Perform POS Tagging (Akkadian)')\n    df_ner   = to_input_target(df_ner_raw,   'phrase', 'ner',           'Do Named Entity Recognition (Akkadian)')\n    df_lemma = to_input_target(df_lemma_raw, 'phrase', 'lemmatization', 'Lemmatize Akkadian text')\n\n    def tokenize(df):\n        dataset = Dataset.from_pandas(df, preserve_index=False)\n\n        def preprocess_function(examples):\n            model_inputs = tokenizer(\n                examples['input'],\n                max_length=MAX_LENGTH,\n                truncation=True\n            )\n            labels = tokenizer(\n                text_target=examples['target'],\n                max_length=MAX_LENGTH,\n                truncation=True\n            )\n            model_inputs['labels'] = labels['input_ids']\n            return model_inputs\n\n        return dataset.map(\n            preprocess_function,\n            batched=True,\n            remove_columns=dataset.column_names\n        )\n\n    tokenized = {\n        'nmt_train': tokenize(df_nmt_train),\n        'nmt_val': tokenize(df_nmt_val),\n        'nmt_add': tokenize(df_nmt_add),\n        'pos': tokenize(df_pos),\n        'ner': tokenize(df_ner),\n        'lemma': tokenize(df_lemma),\n    }\n\n    for name, ds in tokenized.items():\n        print(f\"{name}: {len(ds)} samples\")\n\n    return tokenized\n\n\ndatasets = prepare_all_data(df_ori, df_add, df_multitask)\n\nnmt_dataset = concatenate_datasets(\n    [datasets['nmt_train'], datasets['nmt_add']]\n).shuffle(seed=42)\n# *** OUR CODE ENDS HERE ***","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# *** OUR CODE STARTS HERE ***\n# We mainly wrote the code by ourselves, with the help of Claude Sonnet 4.6 for general debugging purposes\n# Performing Multi-Task Learning\n\n# Training schedule parameters\n# Since we are performing MTL biased towards NMT, auxiliary task weigh less on the loss\n# Such a weight decays over training steps, allowing to progressively concentrate on NMT as the training continues\nAUX_INITIAL_WEIGHT = 0.25\nAUX_DECAY_RATE = 0.0003\nAUX_WARMUP_STEPS = 500\n\n# Batch size has to be a multiple of 4 to allow 1:1:1:1 ration in training batch\n# Here we train for 5 epochs\nBATCH_SIZE = 4\nEPOCHS = 5\nGRADIENT_ACCUM = 4\nOUTPUT_DIR = \"./akkadian-byt5-base-2cpt-mtl5ft-7500-checkpoint\"\nRESUME_CHECPOINT = None\n\n# Helper class for cyclical resampling of auxiliary tasks (created when we did not know if we had enough auxiliary samples to cover all training)\nclass MTLIterableDataset(IterableDataset):\n    def __init__(self, task_datasets):\n        self.task_datasets = task_datasets\n        \n    def __iter__(self):\n        iterators = [\n            (task_id, cycle(ds))\n            for task_id, ds in self.task_datasets\n        ]\n        while True:\n            for task_id, it in iterators:\n                sample = dict(next(it))\n                sample['task_id'] = task_id\n                yield sample\n\n    def __len__(self):\n        min_len = min(len(ds) for _, ds in self.task_datasets)\n        return min_len * len(self.task_datasets)\n\n\nmtl_train_dataset = MTLIterableDataset([\n    (0, nmt_dataset),\n    (1, datasets['pos']),\n    (2, datasets['ner']),\n    (3, datasets['lemma']),\n])\n\n# Normal DataCollator gave us problems because it removed task_id from the dataset, which we needed for distinguishing tasks during training\n# This class overrides this feature\n@dataclass\nclass MTLDataCollator:\n    def __init__(\n        self,\n        base_collator,\n    ):\n        self.base_collator = base_collator\n        \n    def __call__(self, features):\n        features = [dict(f) for f in features]\n        task_ids = [f.pop(\"task_id\", 0) for f in features]\n        batch = self.base_collator(features)\n        batch[\"task_id\"] = torch.tensor(task_ids, dtype=torch.long)\n        return batch\n\ndata_collator = MTLDataCollator(\n    base_collator=DataCollatorForSeq2Seq(\n        tokenizer=tokenizer,\n        model=model,\n        padding=\"longest\"\n    )\n)\n\n# Costum MTL fine-tuner\nclass MTLTrainer(Seq2SeqTrainer):\n    def __init__(\n        self,\n        aux_initial_weight = 0.25,\n        aux_decay_rate = 0.0003,\n        aux_warmup_steps = 500,\n        **kwargs\n    ):\n        super().__init__(**kwargs)\n        self.aux_initial_weight = aux_initial_weight\n        self.aux_decay_rate = aux_decay_rate\n        self.aux_warmup_steps = aux_warmup_steps\n\n    def get_train_dataloader(self):\n        return DataLoader(\n            self.train_dataset,\n            batch_size=BATCH_SIZE,\n            collate_fn=self.data_collator,\n            num_workers=0,\n        )\n\n    # Decay auxiliary weights depending on the training steps (starting from the final warm up step)\n    def _get_aux_weight(self, step):\n        if step < self.aux_warmup_steps:\n            return 0.0\n        return self.aux_initial_weight * np.exp(\n            -self.aux_decay_rate * (step - self.aux_warmup_steps)\n        )\n\n    # The loss is the weighted sum of NTM and scaled auxiliary losses\n    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):\n        current_step = self.state.global_step\n        aux_weight = self._get_aux_weight(current_step)\n    \n        inputs = {k: v for k, v in inputs.items()}\n        task_ids = inputs.pop(\"task_id\")\n    \n        if not model.training:\n            outputs = model(**inputs)\n            loss    = outputs.loss.mean()\n            return (loss, outputs) if return_outputs else loss\n    \n        nmt_mask = task_ids == 0\n        aux_mask = task_ids != 0\n        aux_task_ids = task_ids[aux_mask]\n    \n        nmt_inputs = {k: v[nmt_mask] for k, v in inputs.items()}\n        aux_inputs = {k: v[aux_mask] for k, v in inputs.items()}\n    \n        nmt_outputs = model(**nmt_inputs)\n        nmt_loss = nmt_outputs.loss.mean()\n        total_loss = nmt_loss\n    \n        if aux_weight > 0.0 and aux_mask.any():\n            aux_outputs = model(**aux_inputs)\n            aux_loss_raw = aux_outputs.loss.mean()\n    \n            with torch.no_grad():\n                norm_factor = nmt_loss.detach() / (aux_loss_raw.detach() + 1e-8)\n    \n            aux_loss_normalized = aux_loss_raw * norm_factor\n            total_loss = nmt_loss + aux_weight * aux_loss_normalized\n    \n        self._mtl_nmt_loss_accum = getattr(self, '_mtl_nmt_loss_accum',  0.0) + nmt_loss.item()\n        self._mtl_aux_loss_accum = getattr(self, '_mtl_aux_loss_accum',  0.0) + (aux_loss_raw.item() if aux_weight > 0.0 and aux_mask.any() else 0.0)\n        self._mtl_total_loss_accum = getattr(self, '_mtl_total_loss_accum', 0.0) + total_loss.item()\n        self._mtl_log_steps = getattr(self, '_mtl_log_steps', 0) + 1\n    \n        if current_step % self.args.eval_steps == 0 and current_step > 0:\n            avg_nmt   = self._mtl_nmt_loss_accum   / self._mtl_log_steps\n            avg_aux   = self._mtl_aux_loss_accum   / self._mtl_log_steps\n            avg_total = self._mtl_total_loss_accum / self._mtl_log_steps\n    \n            eval_loss_str = f\"{self._last_eval_loss:.4f}\" if hasattr(self, '_last_eval_loss') else \"...\"\n    \n            print(\n                f\"\\n{'─'*65}\\n\"\n                f\"  Step {current_step:>6}  |  Epoch {self.state.epoch:.2f}\\n\"\n                f\"{'─'*65}\\n\"\n                f\"  {'Train Loss (NMT):':<25} {avg_nmt:.4f}\\n\"\n                f\"  {'Train Loss (Aux):':<25} {avg_aux:.4f}  (weight: {aux_weight:.4f})\\n\"\n                f\"  {'Train Loss (Total):':<25} {avg_total:.4f}\\n\"\n                f\"  {'Eval Loss:':<25} {eval_loss_str}\\n\"\n                f\"{'─'*65}\\n\"\n            )\n    \n            self._mtl_nmt_loss_accum = 0.0\n            self._mtl_aux_loss_accum = 0.0\n            self._mtl_total_loss_accum = 0.0\n            self._mtl_log_steps = 0\n    \n        return (total_loss, nmt_outputs) if return_outputs else total_loss\n\nargs = Seq2SeqTrainingArguments(\n    output_dir=OUTPUT_DIR,\n    eval_strategy=\"steps\",\n    eval_steps=500,\n    learning_rate=5e-4,\n    per_device_train_batch_size=BATCH_SIZE,\n    per_device_eval_batch_size=BATCH_SIZE,\n    weight_decay=0.01,\n    save_strategy=\"epoch\",\n    save_total_limit=1,\n    save_only_model=False, \n    num_train_epochs=EPOCHS,\n    gradient_accumulation_steps=GRADIENT_ACCUM,\n    gradient_checkpointing=True,\n    fp16=False,\n    logging_steps=50,\n    report_to=\"none\",\n    remove_unused_columns=False,\n)\n\ntrainer = MTLTrainer(\n    model=model,\n    args=args,\n    train_dataset=mtl_train_dataset,\n    eval_dataset=datasets['nmt_val'],\n    data_collator=data_collator,\n    aux_initial_weight=AUX_INITIAL_WEIGHT,\n    aux_decay_rate=AUX_DECAY_RATE,\n    aux_warmup_steps=AUX_WARMUP_STEPS,\n)\n\ntrainer.train(resume_from_checkpoint = RESUME_CHECPOINT)\n# *** OUR CODE ENDS HERE ***","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# *** OUR CODE STARTS HERE ***\nsave_path = \"./akkadian-byt5-base-2cpt-mtl5ft-7500\"\n\ntrainer.save_model(save_path)\ntokenizer.save_pretrained(save_path)\n\nzip_path = shutil.make_archive(\"akkadian-byt5-base-2cpt-mtl5ft-7500\", 'zip', save_path)\n# *** OUR CODE ENDS HERE ***","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}