{"cells":[{"cell_type":"markdown","id":"9fa9b3e7","metadata":{},"source":"# Akkadian V10 Inference — Gemma 3n E4B + RAG\n\n**Model**: google/gemma-3n-E4B-it (8B raw, effective 4B)\n**Method**: RAG few-shot — TF-IDF retrieval from training data + LLM translation\n**No fine-tuning needed**: Gemma translates via in-context learning\n\nUpload to Kaggle as notebook → attach model + retriever dataset → Submit"},{"cell_type":"code","execution_count":null,"id":"870ccc6a","metadata":{},"outputs":[],"source":"import os\nimport sys\nimport re\nimport unicodedata\nimport pickle\nimport time\n\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom sklearn.metrics.pairwise import cosine_similarity\n\n# --- Config ---\nSEED = 42\nMAX_NEW_TOKENS = 256\nNUM_EXAMPLES = 5                # few-shot examples per query\nTEMPERATURE = 0.3\nTOP_P = 0.9\n\n# --- Kaggle paths ---\nRETRIEVER_PATH = \"/kaggle/input/v10-retriever/v10_retriever.pkl\"\nCOMP_PATH = \"/kaggle/input/deep-past-initiative-machine-translation\""},{"cell_type":"markdown","id":"27d81806","metadata":{},"source":"## 1. Environment Setup"},{"cell_type":"code","execution_count":null,"id":"6d4b017b","metadata":{},"outputs":[],"source":"import subprocess\nsubprocess.check_call([\n    sys.executable, \"-m\", \"pip\", \"install\", \"-q\",\n    \"transformers>=4.50\",\n    \"accelerate\",\n    \"bitsandbytes\",\n])\n\nimport transformers\nprint(f\"transformers version: {transformers.__version__}\")\n\ntorch.manual_seed(SEED)\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Device: {device}\")\nif device.type == \"cuda\":\n    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n    print(f\"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB\")\n\n\ndef find_path(*candidates):\n    for p in candidates:\n        if os.path.exists(p):\n            return p\n    return candidates[0]\n\n\ncomp_path = find_path(\n    COMP_PATH,\n    \"/kaggle/input/competitions/deep-past-initiative-machine-translation\",\n)\nretriever_path = find_path(\n    RETRIEVER_PATH,\n    \"/kaggle/input/v10-retriever/v10_retriever.pkl\",\n    \"/kaggle/input/akkadian-v10-retriever/v10_retriever.pkl\",\n)\n\nprint(f\"Competition data: {comp_path}\")\nprint(f\"Retriever: {retriever_path}\")"},{"cell_type":"markdown","id":"937a5bf1","metadata":{},"source":"## 2. ASCII Normalization (identical to training)"},{"cell_type":"code","execution_count":null,"id":"dfd656dd","metadata":{},"outputs":[],"source":"_VOWEL_MAP = {\n    '\\u00e0': 'a', '\\u00e1': 'a', '\\u00e2': 'a', '\\u0101': 'a', '\\u00e4': 'a',\n    '\\u00c0': 'A', '\\u00c1': 'A', '\\u00c2': 'A', '\\u0100': 'A', '\\u00c4': 'A',\n    '\\u00e8': 'e', '\\u00e9': 'e', '\\u00ea': 'e', '\\u0113': 'e', '\\u00eb': 'e',\n    '\\u00c8': 'E', '\\u00c9': 'E', '\\u00ca': 'E', '\\u0112': 'E', '\\u00cb': 'E',\n    '\\u00ec': 'i', '\\u00ed': 'i', '\\u00ee': 'i', '\\u012b': 'i', '\\u00ef': 'i',\n    '\\u00cc': 'I', '\\u00cd': 'I', '\\u00ce': 'I', '\\u012a': 'I', '\\u00cf': 'I',\n    '\\u00f2': 'o', '\\u00f3': 'o', '\\u00f4': 'o', '\\u014d': 'o', '\\u00f6': 'o',\n    '\\u00d2': 'O', '\\u00d3': 'O', '\\u00d4': 'O', '\\u014c': 'O', '\\u00d6': 'O',\n    '\\u00f9': 'u', '\\u00fa': 'u', '\\u00fb': 'u', '\\u016b': 'u', '\\u00fc': 'u',\n    '\\u00d9': 'U', '\\u00da': 'U', '\\u00db': 'U', '\\u016a': 'U', '\\u00dc': 'U',\n}\n_CONSONANT_MAP = {\n    '\\u0161': 's', '\\u0160': 'S',\n    '\\u1e63': 's', '\\u1e62': 'S',\n    '\\u1e6d': 't', '\\u1e6c': 'T',\n    '\\u1e2b': 'h', '\\u1e2a': 'H',\n}\n_OCR_MAP = {\n    '\\u201e': '\"', '\\u201c': '\"', '\\u201d': '\"',\n    '\\u2018': \"'\", '\\u2019': \"'\", '\\u201a': \"'\",\n    '\\u02be': \"'\", '\\u02bf': \"'\",\n    '\\u2308': '[', '\\u2309': ']', '\\u230a': '[', '\\u230b': ']',\n}\n_SUBSCRIPT_MAP = str.maketrans({\n    '\\u2080': '0', '\\u2081': '1', '\\u2082': '2', '\\u2083': '3', '\\u2084': '4',\n    '\\u2085': '5', '\\u2086': '6', '\\u2087': '7', '\\u2088': '8', '\\u2089': '9',\n    '\\u2093': 'x',\n})\n_FULL_MAP = str.maketrans({**_VOWEL_MAP, **_CONSONANT_MAP, **_OCR_MAP})\n\n\ndef normalize_transliteration(text) -> str:\n    \"\"\"V4-identical ASCII normalization.\"\"\"\n    if text is None or (isinstance(text, float) and text != text):\n        return \"\"\n    text = str(text)\n    text = unicodedata.normalize(\"NFC\", text)\n    text = text.replace(\"{\", \"(\").replace(\"}\", \")\")\n    text = text.translate(_FULL_MAP)\n    text = text.translate(_SUBSCRIPT_MAP)\n    text = text.replace('\\u2026', ' <gap> ')\n    text = re.sub(r'\\.\\.\\.+', ' <gap> ', text)\n    text = re.sub(r'\\[([^\\]]*)\\]', ' <gap> ', text)\n    text = re.sub(r'\\bx\\b', ' <unk> ', text, flags=re.IGNORECASE)\n    text = re.sub(r'[!?/]', ' ', text)\n    text = re.sub(r'\\s+', ' ', text).strip()\n    return text\n\n\nassert normalize_transliteration(\"šum-ma\") == \"sum-ma\"\nprint(\"Normalization OK\")"},{"cell_type":"markdown","id":"56d261b4","metadata":{},"source":"## 3. Load Gemma 3n E4B-IT (4-bit Quantized)"},{"cell_type":"code","execution_count":null,"id":"6103f4f0","metadata":{},"outputs":[],"source":"from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n\nMODEL_ID = \"google/gemma-3n-E4B-it\"\n\n# Try Kaggle-local model first, fallback to model ID\nmodel_source = find_path(\n    \"/kaggle/input/gemma-3n/transformers/e4b-it/1\",\n    \"/kaggle/input/gemma-3n-e4b-it/1\",\n    MODEL_ID,\n)\nprint(f\"Loading model from: {model_source}\")\n\nbnb_config = BitsAndBytesConfig(\n    load_in_4bit=True,\n    bnb_4bit_compute_dtype=torch.float16,\n    bnb_4bit_quant_type=\"nf4\",\n    bnb_4bit_use_double_quant=True,\n)\n\nt0 = time.time()\ntokenizer = AutoTokenizer.from_pretrained(model_source)\nmodel = AutoModelForCausalLM.from_pretrained(\n    model_source,\n    quantization_config=bnb_config,\n    device_map=\"auto\",\n    torch_dtype=torch.float16,\n)\nprint(f\"Model loaded in {time.time() - t0:.1f}s\")\n\nn_params = sum(p.numel() for p in model.parameters())\nprint(f\"Parameters: {n_params / 1e6:.0f}M\")\n\nif device.type == \"cuda\":\n    mem = torch.cuda.memory_allocated() / 1e9\n    print(f\"GPU memory used: {mem:.2f} GB\")"},{"cell_type":"markdown","id":"ac0f3114","metadata":{},"source":"## 4. Load RAG Retriever"},{"cell_type":"code","execution_count":null,"id":"3fff2392","metadata":{"lines_to_next_cell":1},"outputs":[],"source":"print(f\"\\nLoading retriever from: {retriever_path}\")\nt0 = time.time()\nwith open(retriever_path, \"rb\") as f:\n    retriever_data = pickle.load(f)\n\nvectorizer = retriever_data[\"vectorizer\"]\ntrain_vectors = retriever_data[\"vectors\"]\ntrain_pairs = retriever_data[\"pairs\"]\n\nprint(f\"Loaded in {time.time() - t0:.1f}s\")\nprint(f\"Training pairs: {len(train_pairs):,}\")\nprint(f\"TF-IDF features: {train_vectors.shape[1]:,}\")"},{"cell_type":"markdown","id":"76ef32b9","metadata":{},"source":"## 5. RAG Retrieval Function"},{"cell_type":"code","execution_count":null,"id":"49f41b9b","metadata":{},"outputs":[],"source":"def retrieve_examples(query_src, k=NUM_EXAMPLES):\n    \"\"\"Retrieve top-K most similar training pairs using TF-IDF cosine similarity.\"\"\"\n    q_vec = vectorizer.transform([query_src])\n    sims = cosine_similarity(q_vec, train_vectors)[0]\n    top_k_idx = sims.argsort()[-k:][::-1]\n    return [(train_pairs[i][\"src\"], train_pairs[i][\"tgt\"], sims[i])\n            for i in top_k_idx]\n\n\n# Quick test\n_test_examples = retrieve_examples(\"um-ma ka-ru-um ka-ni-ia-ma\")\nprint(\"\\nRetriever test:\")\nfor s, t, sim in _test_examples[:3]:\n    print(f\"  sim={sim:.3f}: {s[:60]}... → {t[:60]}...\")"},{"cell_type":"markdown","id":"4e9a2033","metadata":{},"source":"## 6. Prompt Construction"},{"cell_type":"code","execution_count":null,"id":"94e73f2a","metadata":{},"outputs":[],"source":"SYSTEM_PROMPT = \"\"\"You are an expert translator of Old Assyrian Akkadian cuneiform texts from approximately 2000 BCE. These are commercial letters, legal documents, and trade records from ancient Mesopotamian merchants.\n\nKey conventions:\n- Transliteration uses hyphens between syllables (e.g., \"um-ma\" = \"thus\")\n- ALL-CAPS words are Sumerograms (logograms): KU.BABBAR = silver, DUMU = son, E.GAL = palace, AN.NA = tin, GIN = shekel, DAM.GAR = merchant\n- <gap> marks damaged/missing text on the tablet\n- Personal names should be kept in their transliterated form\n- Translate naturally into English, preserving the meaning\"\"\"\n\n\ndef build_prompt(test_src, examples):\n    \"\"\"Build a few-shot prompt with retrieved examples.\"\"\"\n    ex_parts = []\n    for i, (src, tgt, sim) in enumerate(examples, 1):\n        # Truncate long examples to keep prompt manageable\n        src_trunc = src[:200] if len(src) > 200 else src\n        tgt_trunc = tgt[:300] if len(tgt) > 300 else tgt\n        ex_parts.append(f\"Example {i}:\\nAkkadian: {src_trunc}\\nEnglish: {tgt_trunc}\")\n\n    examples_str = \"\\n\\n\".join(ex_parts)\n\n    prompt = f\"\"\"{SYSTEM_PROMPT}\n\nHere are similar translation examples for reference:\n\n{examples_str}\n\nNow translate the following Akkadian transliteration to English.\nOutput ONLY the English translation, nothing else.\n\nAkkadian: {test_src}\nEnglish:\"\"\"\n    return prompt"},{"cell_type":"markdown","id":"8224c665","metadata":{},"source":"## 7. Sanity Check"},{"cell_type":"code","execution_count":null,"id":"a77fc16b","metadata":{},"outputs":[],"source":"print(\"\\nSanity check:\")\n_test_src = normalize_transliteration(\"um-ma kà-ru-um kà-ni-ia-ma\")\n_test_examples = retrieve_examples(_test_src, k=3)\n_test_prompt = build_prompt(_test_src, _test_examples)\n\nprint(f\"  Prompt length: {len(_test_prompt)} chars\")\n\n_inputs = tokenizer(_test_prompt, return_tensors=\"pt\").to(model.device)\nprint(f\"  Input tokens: {_inputs.input_ids.shape[1]}\")\n\nwith torch.no_grad():\n    _out = model.generate(\n        **_inputs,\n        max_new_tokens=128,\n        temperature=TEMPERATURE,\n        top_p=TOP_P,\n        do_sample=True,\n    )\n\n_response = tokenizer.decode(_out[0][_inputs.input_ids.shape[1]:],\n                              skip_special_tokens=True).strip()\nprint(f\"  Input: '{_test_src}'\")\nprint(f\"  Output: '{_response[:200]}'\")\n\n# Basic quality check\n_alpha = sum(c.isalpha() and c.isascii() for c in _response) / max(len(_response), 1)\nMODEL_WORKS = len(_response) > 2 and _alpha > 0.4\nprint(f\"  alpha={_alpha:.2f}, MODEL_WORKS={MODEL_WORKS}\")\n\nif not MODEL_WORKS:\n    print(\"  WARNING: Model may not be translating properly!\")"},{"cell_type":"markdown","id":"f809593c","metadata":{},"source":"## 8. Load Test Data"},{"cell_type":"code","execution_count":null,"id":"c0f0fce2","metadata":{},"outputs":[],"source":"test_df = pd.read_csv(os.path.join(comp_path, \"test.csv\"))\nprint(f\"\\nTest samples: {len(test_df)}\")\n\ntest_df[\"src_norm\"] = test_df[\"transliteration\"].apply(normalize_transliteration)\n\nfor i in range(min(5, len(test_df))):\n    print(f\"  [{i}] ({len(test_df['src_norm'].iloc[i])} chars) \"\n          f\"{test_df['src_norm'].iloc[i][:80]}...\")"},{"cell_type":"markdown","id":"57e36e8a","metadata":{},"source":"## 9. Translate All Test Sentences"},{"cell_type":"code","execution_count":null,"id":"f030a929","metadata":{},"outputs":[],"source":"print(\"=\" * 70)\nprint(f\"Gemma 3n E4B + RAG Translation\")\nprint(f\"Test samples: {len(test_df)}\")\nprint(f\"Few-shot examples per query: {NUM_EXAMPLES}\")\nprint(f\"Temperature: {TEMPERATURE}\")\nprint(\"=\" * 70)\n\ntranslations = []\ntotal_time = 0\n\nfor idx, row in test_df.iterrows():\n    src = row[\"src_norm\"]\n    t_start = time.time()\n\n    print(f\"\\n{'='*60}\")\n    print(f\"--- Test [{idx}] / {len(test_df)} ---\")\n    print(f\"Input: {src[:100]}...\")\n\n    # 1. Retrieve similar examples\n    examples = retrieve_examples(src, k=NUM_EXAMPLES)\n    print(f\"Top example (sim={examples[0][2]:.3f}): {examples[0][0][:60]}...\")\n\n    # 2. Build prompt\n    prompt = build_prompt(src, examples)\n\n    # 3. Generate translation\n    inputs = tokenizer(prompt, return_tensors=\"pt\", truncation=True,\n                       max_length=2048).to(model.device)\n\n    with torch.no_grad():\n        output = model.generate(\n            **inputs,\n            max_new_tokens=MAX_NEW_TOKENS,\n            temperature=TEMPERATURE,\n            top_p=TOP_P,\n            do_sample=True,\n        )\n\n    # 4. Extract only the generated part (after prompt)\n    response = tokenizer.decode(\n        output[0][inputs.input_ids.shape[1]:],\n        skip_special_tokens=True\n    ).strip()\n\n    # 5. Clean up response — take only first paragraph/sentence group\n    # Remove any \"Akkadian:\" or \"Example\" if model continues generating\n    for stop_marker in [\"Akkadian:\", \"Example\", \"\\n\\n\\n\", \"---\"]:\n        if stop_marker in response:\n            response = response[:response.index(stop_marker)].strip()\n\n    # 6. Fallback: if response is empty or garbage, use best TM match\n    alpha = sum(c.isalpha() and c.isascii() for c in response) / max(len(response), 1)\n    if len(response) < 3 or alpha < 0.3:\n        # Use best retrieved example's translation as fallback\n        response = examples[0][1]\n        print(f\"  FALLBACK to TM (sim={examples[0][2]:.3f})\")\n\n    t_elapsed = time.time() - t_start\n    total_time += t_elapsed\n\n    print(f\"Translation ({t_elapsed:.1f}s): {response[:120]}\")\n    translations.append(response)\n\n    # Progress estimate\n    if idx > 0:\n        avg_time = total_time / (idx + 1)\n        remaining = avg_time * (len(test_df) - idx - 1)\n        print(f\"  Avg: {avg_time:.1f}s/sample, \"\n              f\"ETA: {remaining / 3600:.1f}h remaining\")\n\nprint(f\"\\n{'='*70}\")\nprint(f\"Total translation time: {total_time / 60:.1f} min\")\nprint(f\"Average per sample: {total_time / max(len(test_df), 1):.1f}s\")"},{"cell_type":"markdown","id":"102f5827","metadata":{},"source":"## 10. Create Submission"},{"cell_type":"code","execution_count":null,"id":"11dc7b7c","metadata":{},"outputs":[],"source":"submission = pd.DataFrame({\n    \"id\": test_df[\"id\"],\n    \"translation\": translations,\n})\n\nassert len(submission) == len(test_df)\nassert submission[\"translation\"].notna().all()\n\nsubmission.to_csv(\"submission.csv\", index=False)\nprint(f\"\\nSubmission saved: submission.csv\")\nprint(f\"Shape: {submission.shape}\")\n\nfor i in range(min(10, len(submission))):\n    print(f\"  [{i}] {submission['translation'].iloc[i][:100]}\")\n\nprint(f\"\\nVerification: {len(submission)} rows, cols={list(submission.columns)}\")\nprint(f\"Any NaN: {submission['translation'].isna().any()}\")\nprint(f\"\\nReady for submission.\")"}],"metadata":{"jupytext":{"cell_metadata_filter":"-all","main_language":"python","notebook_metadata_filter":"-all"}},"nbformat":4,"nbformat_minor":5}