{"cells": [{"cell_type": "code", "execution_count": null, "id": "01860ab3", "metadata": {}, "outputs": [], "source": "\"\"\"\nLoRA Inference for Akkadian Translation (Kaggle)\n\nApplies LoRA adapters to base models and/or weight-averages multiple models\nfor optimized inference using practices from 34.1 LB notebook.\n\nUsage on Kaggle:\n1. Add base model(s) as input datasets\n2. Add LoRA adapter(s) as input datasets (optional)\n3. Add test data from competition\n4. Run notebook - outputs submission.csv\n\nKernel metadata:\n{\n    \"id\": \"manwithacat/lora-inference-v1-0-0\",\n    \"title\": \"lora-inference v1.0.0\",\n    \"code_file\": \"lora-inference-v1-0-0.py\",\n    \"language\": \"python\",\n    \"kernel_type\": \"script\",\n    \"competition\": \"deep-past-initiative-machine-translation\"\n}\n\"\"\""}, {"cell_type": "code", "execution_count": null, "id": "0b981423", "metadata": {}, "outputs": [], "source": "# Install PEFT for LoRA support\nimport subprocess\n\nsubprocess.run([\"pip\", \"install\", \"-q\", \"peft\"])"}, {"cell_type": "code", "execution_count": null, "id": "cb18fb41", "metadata": {}, "outputs": [], "source": "import os\nimport re\nimport gc\nfrom contextlib import nullcontext\n\n# Disable progress bars for cleaner Kaggle logs\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\nos.environ[\"TQDM_DISABLE\"] = \"1\"\nos.environ[\"HF_HUB_DISABLE_PROGRESS_BARS\"] = \"1\"\n\nimport torch\nimport pandas as pd\nimport numpy as np\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM"}, {"cell_type": "code", "execution_count": null, "id": "0f536013", "metadata": {"lines_to_next_cell": 2}, "outputs": [], "source": "# =============================================================================\n# Configuration\n# =============================================================================\n\n\nclass Config:\n    # Paths - adjust based on your Kaggle input datasets\n    TEST_DATA_PATH = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n    OUTPUT_PATH = \"/kaggle/working/submission.csv\"\n\n    # Base model path\n    MODEL_PATHS = [\n        \"/kaggle/input/byt5-base-big-data2\",\n    ]\n    MODEL_WEIGHTS = [1.0]\n\n    # LoRA adapter path\n    LORA_PATHS = [\"/kaggle/input/akkadian-lora-v1-1-0\"]\n\n    # Model settings\n    PREFIX = \"translate Akkadian to English: \"\n    MAX_SOURCE_LENGTH = 512\n\n    # Generation settings (from 34.1 LB notebook)\n    NUM_BEAMS = 6\n    MAX_NEW_TOKENS = 384\n    LENGTH_PENALTY = 1.05\n    EARLY_STOPPING = True\n\n    # Inference settings\n    BATCH_SIZE = 8\n    USE_FP32 = True  # FP32 for score stability (per 34.1 notebook)"}, {"cell_type": "code", "execution_count": null, "id": "32fbcf2f", "metadata": {"lines_to_next_cell": 2}, "outputs": [], "source": "# =============================================================================\n# Preprocessing (from 34.1 LB notebook)\n# =============================================================================\n\n\ndef replace_gaps(text):\n    \"\"\"Normalize gap markers in Akkadian transliteration.\"\"\"\n    if pd.isna(text):\n        return text\n    text = str(text)\n\n    # Large gaps (multiple markers)\n    text = re.sub(\n        r\"\\.3(?:\\s+\\.3)+\\.{3}(?:\\s+\\.{3})+\\s+\\.{3}(?:\\s+\\.{3})+\", \"<big_gap>\", text\n    )\n    text = re.sub(r\"\\.3(?:\\s+\\.3)+\\.{3}(?:\\s+\\.{3})+\", \"<big_gap>\", text)\n    text = re.sub(r\"\\.{3}(?:\\s+\\.{3})+\", \"<big_gap>\", text)\n\n    # Single ellipses\n    text = re.sub(r\"\u2026\u2026\", \"<big_gap>\", text)\n    text = re.sub(r\"\\.\\.\\.\\.\\.\\.\", \"<big_gap>\", text)\n    text = re.sub(r\"\u2026\", \"<big_gap>\", text)\n    text = re.sub(r\"\\.\\.\\.\", \"<big_gap>\", text)\n\n    # Small gaps\n    text = re.sub(r\"xx\", \"<gap>\", text)\n    text = re.sub(r\" x \", \" <gap> \", text)\n\n    return text"}, {"cell_type": "code", "execution_count": null, "id": "5842018c", "metadata": {"lines_to_next_cell": 2}, "outputs": [], "source": "# =============================================================================\n# Weight Averaging\n# =============================================================================\n\n\ndef normalize_weights(weights):\n    \"\"\"Normalize weights to sum to 1.0.\"\"\"\n    total = sum(weights)\n    return [w / total for w in weights]\n\n\ndef weight_average_models(model_paths, weights, device):\n    \"\"\"\n    Load multiple models and average their weights.\n    Returns a single model with averaged parameters.\n    \"\"\"\n    weights = normalize_weights(weights)\n    print(f\"Loading {len(model_paths)} models for ensemble...\")\n    print(f\"Normalized weights: {[f'{w:.4f}' for w in weights]}\")\n\n    # Load all state dicts\n    state_dicts = []\n    for i, path in enumerate(model_paths):\n        print(f\"  Loading model {i + 1}: {path}\")\n        model = AutoModelForSeq2SeqLM.from_pretrained(path)\n        state_dicts.append(model.state_dict())\n        del model\n        gc.collect()\n\n    # Average weights\n    print(\"Averaging weights...\")\n    base_sd = state_dicts[0]\n    averaged_sd = {}\n\n    for key in base_sd:\n        if all(key in sd for sd in state_dicts):\n            averaged_sd[key] = sum(\n                w * sd[key].float() for w, sd in zip(weights, state_dicts)\n            )\n        else:\n            averaged_sd[key] = base_sd[key]\n\n    # Create final model\n    model = AutoModelForSeq2SeqLM.from_pretrained(model_paths[0])\n    model.load_state_dict(averaged_sd)\n\n    return model\n\n\ndef load_model_with_lora(base_path, lora_path):\n    \"\"\"Load base model and merge LoRA adapter.\"\"\"\n    try:\n        from peft import PeftModel\n    except ImportError:\n        raise ImportError(\"PEFT required for LoRA. Install: pip install peft\")\n\n    print(f\"Loading base model: {base_path}\")\n    base_model = AutoModelForSeq2SeqLM.from_pretrained(base_path)\n\n    print(f\"Loading LoRA adapter: {lora_path}\")\n    peft_model = PeftModel.from_pretrained(base_model, lora_path)\n\n    print(\"Merging LoRA into base model...\")\n    merged_model = peft_model.merge_and_unload()\n\n    return merged_model"}, {"cell_type": "code", "execution_count": null, "id": "bcb71e58", "metadata": {"lines_to_next_cell": 2}, "outputs": [], "source": "# =============================================================================\n# Inference Dataset\n# =============================================================================\n\n\nclass InferenceDataset(Dataset):\n    def __init__(self, df, prefix):\n        self.ids = df[\"id\"].tolist()\n        self.texts = [\n            prefix + replace_gaps(str(t)) for t in df[\"transliteration\"].tolist()\n        ]\n\n    def __len__(self):\n        return len(self.texts)\n\n    def __getitem__(self, idx):\n        return self.ids[idx], self.texts[idx]\n\n\ndef collate_fn(batch, tokenizer, max_length):\n    ids, texts = zip(*batch)\n    encoded = tokenizer(\n        list(texts),\n        max_length=max_length,\n        truncation=True,\n        padding=True,\n        return_tensors=\"pt\",\n    )\n    return list(ids), encoded[\"input_ids\"], encoded[\"attention_mask\"]"}, {"cell_type": "code", "execution_count": null, "id": "de8d1a87", "metadata": {}, "outputs": [], "source": "# =============================================================================\n# Main Inference\n# =============================================================================\n\n\ndef main():\n    print(\"=\" * 60)\n    print(\"LoRA INFERENCE FOR AKKADIAN TRANSLATION\")\n    print(\"=\" * 60)\n\n    # Device setup\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\"Device: {device}\")\n    if torch.cuda.is_available():\n        print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n\n    # Load model(s)\n    print(\"\\n\" + \"=\" * 60)\n    print(\"LOADING MODEL\")\n    print(\"=\" * 60)\n\n    if len(Config.MODEL_PATHS) > 1:\n        # Weight averaging multiple models\n        model = weight_average_models(\n            Config.MODEL_PATHS,\n            Config.MODEL_WEIGHTS,\n            device,\n        )\n    else:\n        # Single model\n        model = AutoModelForSeq2SeqLM.from_pretrained(Config.MODEL_PATHS[0])\n\n    # Apply LoRA if specified\n    if Config.LORA_PATHS:\n        for lora_path in Config.LORA_PATHS:\n            print(f\"Applying LoRA: {lora_path}\")\n            try:\n                from peft import PeftModel\n\n                model = PeftModel.from_pretrained(model, lora_path)\n                model = model.merge_and_unload()\n            except ImportError:\n                print(\"WARNING: PEFT not installed, skipping LoRA\")\n\n    # Setup model for inference\n    model.to(device)\n    model.eval()\n\n    # FP32 for stability (per 34.1 notebook)\n    if Config.USE_FP32:\n        model.float()\n        print(\"Using FP32 for inference stability\")\n\n    # Load tokenizer\n    tokenizer = AutoTokenizer.from_pretrained(Config.MODEL_PATHS[0])\n\n    # Load test data\n    print(\"\\n\" + \"=\" * 60)\n    print(\"LOADING TEST DATA\")\n    print(\"=\" * 60)\n\n    test_df = pd.read_csv(Config.TEST_DATA_PATH)\n    print(f\"Test samples: {len(test_df)}\")\n\n    # Create dataloader\n    dataset = InferenceDataset(test_df, Config.PREFIX)\n    dataloader = DataLoader(\n        dataset,\n        batch_size=Config.BATCH_SIZE,\n        shuffle=False,\n        num_workers=0,\n        pin_memory=(device.type == \"cuda\"),\n        collate_fn=lambda batch: collate_fn(batch, tokenizer, Config.MAX_SOURCE_LENGTH),\n    )\n\n    # Inference\n    print(\"\\n\" + \"=\" * 60)\n    print(\"GENERATING TRANSLATIONS\")\n    print(\"=\" * 60)\n\n    all_ids = []\n    all_predictions = []\n\n    torch.set_grad_enabled(False)\n\n    with torch.inference_mode():\n        for batch_idx, (ids, input_ids, attention_mask) in enumerate(dataloader):\n            input_ids = input_ids.to(device)\n            attention_mask = attention_mask.to(device)\n\n            # Generate (FP32 for stability)\n            if Config.USE_FP32:\n                ctx = nullcontext()\n            else:\n                ctx = torch.autocast(device_type=\"cuda\", dtype=torch.float16)\n\n            with ctx:\n                outputs = model.generate(\n                    input_ids=input_ids,\n                    attention_mask=attention_mask,\n                    num_beams=Config.NUM_BEAMS,\n                    max_new_tokens=Config.MAX_NEW_TOKENS,\n                    length_penalty=Config.LENGTH_PENALTY,\n                    early_stopping=Config.EARLY_STOPPING,\n                )\n\n            # Decode\n            decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n            decoded = [d.strip() if d.strip() else \"broken text\" for d in decoded]\n\n            all_ids.extend(ids)\n            all_predictions.extend(decoded)\n\n            if (batch_idx + 1) % 10 == 0:\n                print(f\"  Processed {len(all_ids)}/{len(dataset)} samples\")\n\n    # Create submission\n    print(\"\\n\" + \"=\" * 60)\n    print(\"CREATING SUBMISSION\")\n    print(\"=\" * 60)\n\n    submission = pd.DataFrame(\n        {\n            \"id\": all_ids,\n            \"translation\": all_predictions,\n        }\n    )\n\n    submission.to_csv(Config.OUTPUT_PATH, index=False)\n    print(f\"Saved to: {Config.OUTPUT_PATH}\")\n    print(f\"Total predictions: {len(submission)}\")\n\n    # Show samples\n    print(\"\\n\" + \"=\" * 60)\n    print(\"SAMPLE PREDICTIONS\")\n    print(\"=\" * 60)\n\n    for i in range(min(5, len(submission))):\n        src = test_df.iloc[i][\"transliteration\"]\n        pred = submission.iloc[i][\"translation\"]\n        print(f\"\\n[{i + 1}] {src[:80]}...\")\n        print(f\"    \u2192 {pred[:100]}...\")\n\n    print(\"\\n\" + \"=\" * 60)\n    print(\"INFERENCE COMPLETE\")\n    print(\"=\" * 60)"}, {"cell_type": "code", "execution_count": null, "id": "6ef5384e", "metadata": {}, "outputs": [], "source": "if __name__ == \"__main__\":\n    main()"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}}, "nbformat": 4, "nbformat_minor": 5}