{"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":"competition","sourceId":121150,"databundleVersionId":15833391}],"dockerImageVersionId":31328,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Here's the training code I used. It's pretty standard stuff, nothing more complex than what you'd find in open-source codes. So it's just for an example.\n\nref:\n- https://www.kaggle.com/code/takamichitoda/dpc-starter-train","metadata":{}},{"cell_type":"code","source":"%%writefile run.sh\n#!/bin/bash\n\n# 定义一个运行实验的函数\nrun_experiment() {\n    local TRAIN_DS=$1\n    local VALID_DS=$2\n    local RUN_NAME=$3\n    local RUN_DIR=\"./output/${RUN_NAME}\"\n\n    echo \"=================================================================\"\n    echo \"🚀 开始实验: $RUN_NAME\"\n    echo \"   TRAIN_DS: $TRAIN_DS\"\n    echo \"   VALID_DS: $VALID_DS\"\n    echo \"   RUN_DIR:  $RUN_DIR\"\n    echo \"=================================================================\"\n\n    # 创建文件夹（如果不存在）\n    mkdir -p \"$RUN_DIR\"\n\n    # 启动程序\n    accelerate launch --config_file acc_config.yaml train-byt5.py \\\n        --TRAIN_DS_PATH \"$TRAIN_DS\" \\\n        --VALID_DS_PATH \"$VALID_DS\" \\\n        --OUTPUT_DIR \"$RUN_DIR\" \\\n        --RUN_NAME \"$RUN_NAME\" 2>&1 | tee \"$RUN_DIR/my_log.txt\"\n\n    echo \"✅ 实验 $RUN_NAME 已完成！日志路径: $RUN_DIR/my_log.txt\"\n    echo \"\"\n}\n\nrun_experiment \\\n    \"/2022533109/guguoqin/deep-past/datasets/data1/all.json\" \\\n    \"None\" \\\n    \"byt5-xl/hf-20260323-data1-all\"\n\nrun_experiment \\\n    \"/2022533109/guguoqin/deep-past/datasets/data1_llmlabel/all.json\" \\\n    \"None\" \\\n    \"byt5-xl/hf-20260323-data1_llmlabel-all\"\n\nrun_experiment \\\n    \"/2022533109/guguoqin/deep-past/datasets/data2_llmlabel_synth12/all.json\" \\\n    \"None\" \\\n    \"byt5-xl/hf-20260323-data2_llmlabel_synth12-all\"\n\nrun_experiment \\\n    \"/2022533109/guguoqin/deep-past/datasets/data3_llmlabel_synth12/all.json\" \\\n    \"None\" \\\n    \"byt5-xl/hf-20260323-data3_llmlabel_synth12-all\"\n\necho \"🎉 所有实验运行完毕！\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-24T06:39:31.288974Z","iopub.execute_input":"2026-03-24T06:39:31.289478Z","iopub.status.idle":"2026-03-24T06:39:31.298052Z","shell.execute_reply.started":"2026-03-24T06:39:31.289435Z","shell.execute_reply":"2026-03-24T06:39:31.297041Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile train-byt5.py\nimport argparse\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--TRAIN_DS_PATH\", type=str, default=\"/2022533109/guguoqin/deep-past/datasets/data2_llmlabel_synth2/all.json\")\nparser.add_argument(\"--VALID_DS_PATH\", type=str, default=None)\nparser.add_argument(\"--OUTPUT_DIR\", type=str, default=\"./output/byt5-xl/hf-20260321-data2_llmlabel_synth2-all\")\nparser.add_argument(\"--RUN_NAME\", type=str, default=\"byt5-xl/hf-20260321-data2_llmlabel_synth2-all\")\nargs = parser.parse_args()\nif args.VALID_DS_PATH == \"None\":\n    args.VALID_DS_PATH = None\n\nclass Config:\n    # Train\n    PREFIX_EN = \"Translate Akkadian to English: \"\n    PREFIX_TR = \"Translate Akkadian to Turkish: \"\n    PREFIX_DE = \"Translate Akkadian to German: \"\n    TRAIN_MAX_LENGTH = 1024\n    \n    EPOCHS = 3\n    LEARNING_RATE = 2e-4\n    MIN_LEARNING_RATE = 1.33e-4 # 0.66e-4\n    BATCH_SIZE = 8\n    GRADIENT_ACCUMULATION_STEPS = 1\n\n    # Basic settings\n    CUDA_VISIBLE_DEVICES = \"0,1,2,3,4,5,6,7\"\n    MODEL_NAME = \"/2022533109/guguoqin/models/google/byt5-xl\"\n\n    TRAIN_DS_PATH = args.TRAIN_DS_PATH\n    VALID_DS_PATH = args.VALID_DS_PATH\n\n    OUTPUT_DIR = args.OUTPUT_DIR\n    RUN_NAME = args.RUN_NAME\n\n    # Inference\n    DO_INFER = False\n    INFER_MAX_LENGTH = 1024\n    INFER_NUM_BEAMS = 4\n    INFER_MAX_NEW_TOKENS = 512\n    INFER_BATCH_SIZE = 16\n    TARGET_CKPT = 3\n\nimport os\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = Config.CUDA_VISIBLE_DEVICES\nos.environ[\"SWANLAB_PROJECT\"]=\"deep-past\"\nos.environ[\"SWANLAB_LOG_DIR\"]=Config.OUTPUT_DIR\nimport gc\nimport re\nimport math\nimport json\nimport random\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom sklearn.model_selection import train_test_split\nfrom datasets import Dataset\nfrom transformers import (\n    AutoTokenizer, \n    AutoModelForSeq2SeqLM, \n    DataCollatorForSeq2Seq, \n    Seq2SeqTrainingArguments, \n    Seq2SeqTrainer\n)\nimport evaluate\n\ndef seed_everything(seed=42):\n    torch.manual_seed(seed)\n    np.random.seed(seed)\n    torch.cuda.manual_seed(seed)\nseed_everything()\n\n# ==========================================\n# 1. Load data\n# ==========================================\ndef load_dataset(file_path):\n    with open(file_path, 'r', encoding='utf-8') as f:\n        data = json.load(f)\n    df = []\n    for item in data:\n        df.append({\n            \"uuid\": item[\"uuid\"],\n            \"publication_name\": item['publication_name'],\n            \"chapter_name\": item['chapter_name'],\n            \"sentence_id\": item.get(\"sentence_id\", 1),\n            \"target_lang_code\": item[\"messages\"][0][\"content\"][0][\"target_lang_code\"],\n            \"transliteration\": item[\"messages\"][0][\"content\"][0][\"text\"],\n            \"translation\": item[\"messages\"][1][\"content\"]\n        })\n    df = pd.DataFrame(df)\n    return df\n\ndef get_token_limits(transliteration_len):\n    # 根据你的拟合结果填入系数\n    m, c = 1.1427, -1.1173\n    up, down = 200, 200\n    \n    max_tokens = int(m * transliteration_len + (c + up))\n    min_tokens = max(0, int(m * transliteration_len + (c - down)))\n    \n    return min_tokens, max_tokens\n\ntrain_df = load_dataset(Config.TRAIN_DS_PATH)\nprint(f\"Load Train Data: {Config.TRAIN_DS_PATH}\")\ntrain_df['transliteration_char_length'] = train_df.apply(lambda row: len(row['transliteration']), axis=1)\ntrain_df['translation_char_length'] = train_df.apply(lambda row: len(row['translation']), axis=1)\nprint(f\"Original Train Data: {len(train_df)} docs\")\nprint(f\"Desc of transliteration_char_length: {train_df['transliteration_char_length'].describe()}\")\nprint(f\"Desc of translation_char_length: {train_df['translation_char_length'].describe()}\")\nif Config.VALID_DS_PATH:\n    valid_df = load_dataset(Config.VALID_DS_PATH)\n    valid_df['transliteration_char_length'] = valid_df.apply(lambda row: len(row['transliteration']), axis=1)\n    valid_df['translation_char_length'] = valid_df.apply(lambda row: len(row['translation']), axis=1)\n    valid_df = valid_df.sort_values(by='transliteration_char_length', ascending=False).reset_index(drop=True)\n    valid_df[['min_t', 'max_t']] = valid_df.apply(\n        lambda row: get_token_limits(row['transliteration_char_length']), \n        axis=1, result_type='expand'\n    )\n    print(f\"Original Valid Data: {len(valid_df)} docs\")\n    print(f\"Desc of transliteration_char_length: {valid_df['transliteration_char_length'].describe()}\")\n    print(f\"Desc of translation_char_length: {valid_df['translation_char_length'].describe()}\")\n\n# Convert to Hugging Face Dataset format\ntrain_dataset = Dataset.from_pandas(train_df)\nif Config.VALID_DS_PATH:\n    valid_dataset = Dataset.from_pandas(valid_df)\n    effective_batch_size = Config.BATCH_SIZE * Config.GRADIENT_ACCUMULATION_STEPS\n    steps_per_epoch = math.ceil(len(train_dataset) / effective_batch_size)\n    target_step = steps_per_epoch * Config.TARGET_CKPT\n    dynamic_ckpt_path = os.path.join(Config.OUTPUT_DIR, f\"checkpoint-{target_step}\")\n    print(f\"Calculated steps per epoch: {steps_per_epoch}\")\n    print(f\"Targeting INFER CKPT {Config.TARGET_CKPT} -> Will loading checkpoint from: {dynamic_ckpt_path}\")\n\n# ==========================================\n# 2. Tokenization & preprocessing\n# ==========================================\ntokenizer = AutoTokenizer.from_pretrained(Config.MODEL_NAME)\n\ndef preprocess_function(examples):\n    inputs = []\n    for lang_code, ex in zip(examples[\"target_lang_code\"], examples[\"transliteration\"]):\n        if lang_code == \"en\":\n            inputs.append(Config.PREFIX_EN + str(ex))\n        elif lang_code == \"tr\":\n            inputs.append(Config.PREFIX_TR + str(ex))\n        elif lang_code == \"de\":\n            inputs.append(Config.PREFIX_DE + str(ex))\n    targets = [str(ex) for ex in examples[\"translation\"]]\n    \n    model_inputs = tokenizer(inputs, max_length=Config.TRAIN_MAX_LENGTH, truncation=True)\n    labels = tokenizer(targets, max_length=Config.TRAIN_MAX_LENGTH, truncation=True)\n    \n    model_inputs[\"labels\"] = labels[\"input_ids\"]\n    return model_inputs\n\ntokenized_train = train_dataset.map(preprocess_function, batched=True)\nif Config.VALID_DS_PATH:\n    tokenized_valid = valid_dataset.map(preprocess_function, batched=True)\n\n# ==========================================\n# 3. Model training\n# ==========================================\ngc.collect()\ntorch.cuda.empty_cache()\nmodel = AutoModelForSeq2SeqLM.from_pretrained(Config.MODEL_NAME)\ndata_collator = DataCollatorForSeq2Seq(\n    tokenizer, \n    model=model\n)\n\n# Metric\nchrf_metric = evaluate.load(\"evaluate/metrics/chrf/chrf.py\")\nsacrebleu_metric = evaluate.load(\"evaluate/metrics/sacrebleu/sacrebleu.py\")\ndef compute_metrics(eval_preds):\n    preds, labels = eval_preds\n    if isinstance(preds, tuple): \n        preds = preds[0]\n    if hasattr(preds, \"ndim\") and preds.ndim == 3:\n        preds = np.argmax(preds, axis=-1)\n    \n    preds = preds.astype(np.int64)\n    preds = np.where(preds < 0, tokenizer.pad_token_id, preds)\n    preds = np.clip(preds, 0, tokenizer.vocab_size - 1)\n    \n    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)\n    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n    \n    decoded_preds = [pred.strip() for pred in decoded_preds]\n    decoded_labels = [[label.strip()] for label in decoded_labels]\n    \n    chrf_result = chrf_metric.compute(predictions=decoded_preds, references=decoded_labels, word_order=2)\n    chrf_score = chrf_result[\"score\"]\n    \n    bleu_result = sacrebleu_metric.compute(predictions=decoded_preds, references=decoded_labels)\n    bleu_score = bleu_result[\"score\"]\n    \n    geo_mean = (chrf_score * bleu_score) ** 0.5 if chrf_score > 0 and bleu_score > 0 else 0.0\n    \n    return {\n        \"chrf\": chrf_score,\n        \"bleu\": bleu_score,\n        \"geo_mean\": geo_mean\n    }\n\nclass FocalLossCallback:\n    def __init__(self, gamma=1.0, alpha=1.0, ignore_index=-100):\n        \"\"\"\n        初始化 Focal Loss\n        :param gamma: 聚焦参数，越大越关注难分类的样本 (通常取 2.0)\n        :param alpha: 平衡参数 (通常取 1.0，除非你有明确的类别权重需求)\n        :param ignore_index: 忽略的 Label 索引 (HuggingFace 默认为 -100)\n        \"\"\"\n        self.gamma = gamma\n        self.alpha = alpha\n        self.ignore_index = ignore_index\n\n    def __call__(self, outputs, labels, num_items_in_batch=None):\n        \"\"\"\n        符合 HF Trainer compute_loss_func 签名的调用函数\n        \"\"\"\n        # 1. 维度展平 (Flatten)\n        # Logits: [batch_size, seq_len, vocab_size] -> [N, vocab_size]\n        # Labels: [batch_size, seq_len]             -> [N]\n        logits = outputs.get(\"logits\") if isinstance(outputs, dict) else outputs[0]\n        logits = logits.view(-1, logits.size(-1))\n        labels = labels.view(-1)\n\n        # 2. 计算基础 Cross Entropy (不进行 reduce，保留每个 Token 的 loss)\n        ce_loss = F.cross_entropy(\n            logits, \n            labels, \n            reduction='none', \n            ignore_index=self.ignore_index\n        )\n\n        # 3. 计算概率 pt\n        # ce_loss = -log(pt) => pt = exp(-ce_loss)\n        pt = torch.exp(-ce_loss)\n\n        # 4. 计算 Focal Loss: - alpha * (1 - pt)^gamma * log(pt)\n        loss = self.alpha * (1 - pt) ** self.gamma * ce_loss\n\n        # 5. 计算平均值\n        if self.ignore_index is not None:\n            active_elements_mask = labels != self.ignore_index\n            return loss.sum() / active_elements_mask.sum()\n        else:\n            return loss.mean()\n\nargs = Seq2SeqTrainingArguments(\n    # seed=Config.SEED,\n    # data_seed=Config.SEED,\n    num_train_epochs=Config.EPOCHS,\n    learning_rate=Config.LEARNING_RATE,\n    lr_scheduler_type=\"polynomial\",\n    lr_scheduler_kwargs={\"lr_end\": Config.MIN_LEARNING_RATE, \"power\": 1.0},\n    weight_decay=0.01,\n    # optim=\"adafactor\",  # adamw_torch\n    # label_smoothing_factor=0.1,\n    \n    # === Key fixes ===\n    fp16=False,  \n    tf32=True,\n    group_by_length=False,\n    per_device_train_batch_size=Config.BATCH_SIZE,\n    per_device_eval_batch_size=Config.INFER_BATCH_SIZE,\n    gradient_checkpointing=True,\n    gradient_accumulation_steps=Config.GRADIENT_ACCUMULATION_STEPS, \n    generation_max_length=Config.INFER_MAX_NEW_TOKENS, \n    generation_num_beams=Config.INFER_NUM_BEAMS,  \n    # ======================\n\n    output_dir=Config.OUTPUT_DIR,\n    eval_strategy=\"epoch\" if Config.VALID_DS_PATH else \"no\",\n    save_strategy=\"epoch\",\n    # save_total_limit=8,\n    save_only_model=True,\n    predict_with_generate=True,\n    logging_steps=10,  \n    report_to=\"swanlab\",\n    run_name=Config.RUN_NAME, \n)\n\ntrainer = Seq2SeqTrainer(\n    model=model,\n    args=args,\n    train_dataset=tokenized_train,\n    eval_dataset=tokenized_valid if Config.VALID_DS_PATH else None,\n    data_collator=data_collator,\n    tokenizer=tokenizer,\n    compute_metrics=compute_metrics,\n    # compute_loss_func=FocalLossCallback(gamma=1.0, alpha=1.0, ignore_index=-100)\n)\n\nprint(\"Starting Training (FP32 mode)...\")\ntrainer.train()\n\n# ==========================================\n# 4. Model inference\n# ==========================================\nif Config.DO_INFER and Config.VALID_DS_PATH:\n    from torch.utils.data import Dataset, DataLoader\n\n    class AkkadDataset(Dataset):\n        def __init__(self, df):\n            self.ids = df['uuid'].tolist()\n            self.texts = []\n            for text, lang in zip(df['transliteration'], df['target_lang_code']):\n                if lang == \"en\":\n                    self.texts.append(Config.PREFIX_EN + str(text))\n                elif lang == \"tr\":\n                    self.texts.append(Config.PREFIX_TR + str(text))\n                elif lang == \"de\":\n                    self.texts.append(Config.PREFIX_DE + str(text))\n                else:\n                    self.texts.append(Config.PREFIX_EN + str(text)) # 默认 fallback\n            self.min_tokens = df['min_t'].tolist()\n            self.max_tokens = df['max_t'].tolist()\n\n        def __len__(self): return len(self.ids)\n\n        def __getitem__(self, i):\n            return {\n                \"id\": self.ids[i],\n                \"text\": self.texts[i],\n                \"min_t\": self.min_tokens[i],\n                \"max_t\": self.max_tokens[i]\n            }\n\n    def collate_fn(batch):\n        ids = [x[\"id\"] for x in batch]\n        texts = [x[\"text\"] for x in batch]\n        batch_min = int(np.min([x[\"min_t\"] for x in batch]))\n        batch_max = int(np.max([x[\"max_t\"] for x in batch]))\n        \n        inputs = tokenizer(texts, max_length=Config.INFER_MAX_LENGTH, padding=True, truncation=True, return_tensors=\"pt\")\n        \n        return ids, inputs, batch_min, batch_max\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model = AutoModelForSeq2SeqLM.from_pretrained(dynamic_ckpt_path, dtype=torch.float32).to(device).eval()\n    tokenizer = AutoTokenizer.from_pretrained(dynamic_ckpt_path)\n\n    loader = DataLoader(\n        AkkadDataset(valid_df), \n        batch_size=Config.INFER_BATCH_SIZE, \n        shuffle=False, \n        collate_fn=collate_fn\n    )\n\n    results = []\n\n    with torch.inference_mode():\n        for ids, inputs, min_t, max_t in tqdm(loader, desc=\"Inference\"):\n            \n            generate_kwargs = {\n                \"num_beams\": Config.INFER_NUM_BEAMS,\n                \"do_sample\": False,\n                \"early_stopping\": True,\n                \"max_new_tokens\": Config.INFER_MAX_NEW_TOKENS, # max_t,\n                # \"min_new_tokens\": min_t,\n                \"length_penalty\": 1.05,\n                # \"repetition_penalty\": 1.0\n            }\n\n            outputs = model.generate(\n                input_ids=inputs.input_ids.to(device),\n                attention_mask=inputs.attention_mask.to(device),\n                **generate_kwargs\n            )\n            \n            decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n            # cleaned = [hnr(txt) for txt in decoded]\n            \n            results.extend(zip(ids, decoded))\n\n    submission = pd.DataFrame(results, columns=['uuid', 'translation'])\n    submission = submission.sort_values(by='uuid', ascending=True).reset_index(drop=True)\n    submission.to_csv(os.path.join(Config.OUTPUT_DIR, 'submission-train.csv'), index=False, encoding='utf-8')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-03-24T06:51:42.927431Z","iopub.execute_input":"2026-03-24T06:51:42.927792Z","iopub.status.idle":"2026-03-24T06:51:42.941506Z","shell.execute_reply.started":"2026-03-24T06:51:42.927760Z","shell.execute_reply":"2026-03-24T06:51:42.940577Z"}},"outputs":[],"execution_count":null}]}