{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🧠 IOL-AI 2026 — Solution Notebook\n",
    "\n",
    "**International Linguistics Olympiad AI Challenge**  \n",
    "Can a system reason out a language it has never seen?\n",
    "\n",
    "---\n",
    "\n",
    "## Strategy\n",
    "This notebook has **two modes**:\n",
    "\n",
    "| Mode | When to use | Model |\n",
    "|------|-------------|-------|\n",
    "| **API mode** (Section 2) | Local / Colab prototyping | Gemini 2.0 Flash (free tier) |\n",
    "| **Local model mode** (Section 3) | HuggingFace sandbox / final submission | Qwen2.5-7B-Instruct (4-bit) |\n",
    "\n",
    "### Scoring formula\n",
    "$$\\text{score} = \\sqrt{\\text{EM}_w \\times \\text{chrF}_w}$$\n",
    "\n",
    "We maximise both exact matches AND near-miss character overlap."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 0 — Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ Dependencies installed\n"
     ]
    }
   ],
   "source": [
    "# Install required packages\n",
    "!pip install -q google-generativeai pandas sacrebleu tqdm\n",
    "\n",
    "print('✓ Dependencies installed')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 1 — Data Setup & Mock Data\n",
    "\n",
    "The real test CSV is mounted at `/tmp/data/test.csv` on the HF eval sandbox.  \n",
    "For local development we create a **mock dataset** with real IOL-style problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mode: MOCK (local)\n",
      "Data: mock_test.csv\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "import re\n",
    "import time\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "\n",
    "# ── Path selection ─────────────────────────────────────────────────────────\n",
    "HF_DATA_PATH = '/tmp/data/test.csv'\n",
    "LOCAL_DATA_PATH = 'mock_test.csv'\n",
    "\n",
    "USE_MOCK = not os.path.exists(HF_DATA_PATH)\n",
    "DATA_PATH = HF_DATA_PATH if not USE_MOCK else LOCAL_DATA_PATH\n",
    "\n",
    "print(f'Mode: {\"MOCK (local)\" if USE_MOCK else \"REAL (HF sandbox)\"}')\n",
    "print(f'Data: {DATA_PATH}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ Created mock dataset with 3 problems\n"
     ]
    }
   ],
   "source": [
    "# ── Create realistic mock data (IOL-style problems) ──────────────────────────\n",
    "MOCK_PROBLEMS = [\n",
    "    {\n",
    "        'id': '012023020100',\n",
    "        'context': '''Here are some sentences in Hakhun Tangsa (a Tibeto-Burman language of India) with English translations:\n",
    "ŋa ka kɤ ne | Do I go?\n",
    "nɤ ʒip tuʔ ne | Did you sleep?\n",
    "ŋa ʒip kɤ ne | Do I sleep?\n",
    "nɤ ka tuʔ ne | Did you go?\n",
    "ŋa man kɤ ne | Do I eat?\n",
    "nɤ man tuʔ ne | Did you eat?\n",
    "ŋa ʒip ku ne | Does/Will I sleep?\n",
    "nɤ ka ku ne | Does/Will you go?''',\n",
    "        'query': 'Translate into English:\\n1. nɤ ʒip ku ne\\n2. ŋa man tuʔ ne\\n3. nɤ ka kɤ ne',\n",
    "        'work_lang': 'eng_Latn',\n",
    "        'task_lang': 'tcz_Latn',\n",
    "        'task_type': 'translation',\n",
    "        'eval_type': 'single',\n",
    "    },\n",
    "    {\n",
    "        'id': '012023030200',\n",
    "        'context': '''Here are the squares of 1 through 10 in Ndom (a language of Papua New Guinea), given in arbitrary order:\n",
    "mer an thef abo thonith  [= 25]\n",
    "nif thef abo tondor abo mer abo thonith  [= 64]\n",
    "tondor abo mer abo thonith  [= 49]\n",
    "thonith  [= 1]\n",
    "mer abo thonith  [= 9]\n",
    "ithin  [= 100]\n",
    "nif abo thonith  [= 16]\n",
    "\n",
    "Ndom number system (base 6):\n",
    "thonith = 1\n",
    "mer = 2\n",
    "tondor = 3 (also used as 3-groups)\n",
    "thef = 6\n",
    "nif = 36\n",
    "ithin = 216\n",
    "\n",
    "Simpler data — Ndom numbers:\n",
    "thonith = 1\n",
    "mer = 2  \n",
    "tondor = 3\n",
    "an = 4\n",
    "thef = 5\n",
    "nif = 6\n",
    "nif abo thonith = 7\n",
    "nif abo mer = 8\n",
    "nif abo tondor = 9\n",
    "nif abo an = 10\n",
    "nif abo thef = 11\n",
    "mer nif = 12\n",
    "ithin = 36\n",
    "ithin abo thonith = 37''',\n",
    "        'query': 'Write in numerals:\\n1. nif abo thonith\\n2. mer nif abo an\\n3. ithin abo nif abo mer',\n",
    "        'work_lang': 'eng_Latn',\n",
    "        'task_lang': 'ndo_Latn',\n",
    "        'task_type': 'text_to_num',\n",
    "        'eval_type': 'single',\n",
    "    },\n",
    "    {\n",
    "        'id': '012024010300',\n",
    "        'context': '''Here are some words in Apurinã (an Arawakan language of Brazil) with English translations:\n",
    "kama = macaw (a bird)\n",
    "ãkiti = monkey\n",
    "tsura = jaguar\n",
    "kamana = macaws\n",
    "ãkitia = monkeys\n",
    "tsuraa = jaguars\n",
    "apa = father\n",
    "apana = fathers\n",
    "ama = mother\n",
    "amaa = mothers\n",
    "putsu = stone\n",
    "putsuna = stones''',\n",
    "        'query': 'Fill in the blanks:\\n1. kite = fish → fishes = ____\\n2. nãka = tree → trees = ____\\n3. ãkitia = monkeys → monkey = ____',\n",
    "        'work_lang': 'eng_Latn',\n",
    "        'task_lang': 'apu_Latn',\n",
    "        'task_type': 'fill_blanks',\n",
    "        'eval_type': 'single',\n",
    "    },\n",
    "]\n",
    "\n",
    "if USE_MOCK:\n",
    "    mock_df = pd.DataFrame(MOCK_PROBLEMS)\n",
    "    mock_df.to_csv(LOCAL_DATA_PATH, index=False)\n",
    "    print(f'✓ Created mock dataset with {len(mock_df)} problems')\n",
    "    mock_df[['id','task_type','task_lang']].head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 3 problems\n",
      "\n",
      "Columns: ['id', 'context', 'query', 'work_lang', 'task_lang', 'task_type', 'eval_type']\n",
      "\n",
      "Task type distribution:\n",
      "task_type\n",
      "translation    1\n",
      "text_to_num    1\n",
      "fill_blanks    1\n",
      "\n",
      "--- First problem preview ---\n",
      "Here are some sentences in Hakhun Tangsa (a Tibeto-Burman language of India) with English translations:\n",
      "ŋa ka kɤ ne | Do I go?\n",
      "nɤ ʒip tuʔ ne | Did you sleep?\n",
      "ŋa ʒip kɤ ne | Do I sleep?\n",
      "nɤ ka tuʔ ne | Did you go?\n",
      "ŋa man kɤ ne | Do I eat?\n",
      "nɤ man tuʔ ne | Did you eat?\n",
      "ŋa ʒip ku ne | Does/Will I sleep?\n",
      "nɤ ka ku ne | Does/Will you go?\n",
      "...\n",
      "Translate into English:\n",
      "1. nɤ ʒip ku ne\n",
      "2. ŋa man tuʔ ne\n",
      "3. nɤ ka kɤ ne\n"
     ]
    }
   ],
   "source": [
    "# ── Load and inspect the data ──────────────────────────────────────────────\n",
    "df = pd.read_csv(DATA_PATH, dtype=str).fillna('')\n",
    "print(f'Loaded {len(df)} problems')\n",
    "print()\n",
    "print('Columns:', df.columns.tolist())\n",
    "print()\n",
    "print('Task type distribution:')\n",
    "print(df['task_type'].value_counts().to_string())\n",
    "print()\n",
    "print('--- First problem preview ---')\n",
    "print(df.iloc[0]['context'][:500])\n",
    "print('...')\n",
    "print(df.iloc[0]['query'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2 — API Mode: Gemini 2.0 Flash (Prototyping)\n",
    "\n",
    "Use this in Colab to quickly prototype and iterate. **Does NOT work in the HF sandbox** (no internet).  \n",
    "Set your Gemini API key from [aistudio.google.com](https://aistudio.google.com)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gemini API: DISABLED (no key set)\n"
     ]
    }
   ],
   "source": [
    "# ═══════════════════════════════════════════════════════════════════\n",
    "# 🔑  SET YOUR GEMINI API KEY HERE  (or use Colab Secrets sidebar)\n",
    "# ═══════════════════════════════════════════════════════════════════\n",
    "GEMINI_API_KEY = ''  # <-- paste your key, or leave blank to skip\n",
    "\n",
    "# Alternatively, load from Colab secrets:\n",
    "try:\n",
    "    from google.colab import userdata\n",
    "    GEMINI_API_KEY = userdata.get('GEMINI_API_KEY') or GEMINI_API_KEY\n",
    "except Exception:\n",
    "    pass\n",
    "\n",
    "USE_API = bool(GEMINI_API_KEY)\n",
    "print(f'Gemini API: {\"ENABLED\" if USE_API else \"DISABLED (no key set)\"}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ Prompt utilities ready\n"
     ]
    }
   ],
   "source": [
    "# ── Prompt engineering ─────────────────────────────────────────────────────\n",
    "\n",
    "SYSTEM_PROMPT = \"\"\"You are an expert solver of International Linguistics Olympiad (IOL) problems.\n",
    "These puzzles are COMPLETELY self-contained: all needed information is in the problem text.\n",
    "NEVER use external language knowledge — deduce everything from the examples given.\n",
    "\n",
    "## Method (follow every time):\n",
    "1. READ the full problem carefully.\n",
    "2. IDENTIFY patterns: morphemes, grammatical markers, number systems, word order, affixes.\n",
    "3. BUILD a concise rule table / lexicon from the examples.\n",
    "4. VERIFY your rules on a few given examples before answering.\n",
    "5. APPLY rules systematically to each numbered item.\n",
    "\n",
    "## Output format (CRITICAL):\n",
    "First, show your reasoning (2-10 lines). Then output:\n",
    "\n",
    "===ANSWERS===\n",
    "[answer to item 1]\n",
    "[answer to item 2]\n",
    "...\n",
    "\n",
    "One line per numbered item, in order, NO numbering, NO extra text.\n",
    "If multiple answers are valid, write them all on one line separated by \" / \".\"\"\"\n",
    "\n",
    "\n",
    "def build_user_prompt(context: str, query: str, task_type: str = '') -> str:\n",
    "    hints = {\n",
    "        'text_to_num': '\\n[TASK TYPE: Convert written number words → Arabic numerals. Map each number word methodically.]\\n',\n",
    "        'num_to_text': '\\n[TASK TYPE: Convert Arabic numerals → number words in the target language.]\\n',\n",
    "        'fill_blanks': '\\n[TASK TYPE: Fill in missing words using grammatical patterns from the examples.]\\n',\n",
    "        'match_letters': '\\n[TASK TYPE: Map each symbol/letter to its romanised equivalent using the given correspondences.]\\n',\n",
    "        'translation': '\\n[TASK TYPE: Translate sentences using the vocabulary and grammar patterns shown.]\\n',\n",
    "    }\n",
    "    hint = hints.get(task_type, '')\n",
    "    return f'{context.strip()}\\n{hint}\\n{query.strip()}'\n",
    "\n",
    "\n",
    "def count_items(query: str) -> int:\n",
    "    \"\"\"Count numbered items (1., 2., 17., …) in the query.\"\"\"\n",
    "    return max(len(re.findall(r'^\\s*\\d+\\.', query, re.MULTILINE)), 1)\n",
    "\n",
    "\n",
    "def parse_answers(raw: str, n_items: int) -> list:\n",
    "    \"\"\"Extract ===ANSWERS=== block, or fall back to last N lines.\"\"\"\n",
    "    if '===ANSWERS===' in raw:\n",
    "        block = raw.split('===ANSWERS===', 1)[1].strip()\n",
    "        lines = [l.strip() for l in block.splitlines() if l.strip()]\n",
    "        if lines:\n",
    "            return lines[:n_items] if len(lines) >= n_items else lines\n",
    "    # Fallback\n",
    "    lines = [l.strip() for l in raw.splitlines() if l.strip()]\n",
    "    return lines[-n_items:] if len(lines) >= n_items else lines\n",
    "\n",
    "\n",
    "def extract_explanation(raw: str) -> str:\n",
    "    \"\"\"Extract the reasoning section (before ===ANSWERS===) as explanation.\"\"\"\n",
    "    if '===ANSWERS===' in raw:\n",
    "        reasoning = raw.split('===ANSWERS===')[0].strip()\n",
    "    else:\n",
    "        reasoning = raw\n",
    "    # Trim to 500 chars for the explanation column\n",
    "    return reasoning[:500]\n",
    "\n",
    "\n",
    "print('✓ Prompt utilities ready')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Gemini API inference ────────────────────────────────────────────────────\n",
    "if USE_API:\n",
    "    import google.generativeai as genai\n",
    "\n",
    "    genai.configure(api_key=GEMINI_API_KEY)\n",
    "\n",
    "    generation_config = genai.GenerationConfig(\n",
    "        temperature=0.0,       # Deterministic\n",
    "        max_output_tokens=1024,\n",
    "    )\n",
    "\n",
    "    gemini_model = genai.GenerativeModel(\n",
    "        model_name='gemini-2.0-flash',\n",
    "        system_instruction=SYSTEM_PROMPT,\n",
    "        generation_config=generation_config,\n",
    "    )\n",
    "\n",
    "    print('✓ Gemini model initialised: gemini-2.0-flash')\n",
    "\n",
    "    # Quick test\n",
    "    test_resp = gemini_model.generate_content('Say: ready')\n",
    "    print('Test:', test_resp.text.strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Run API inference over the dataset ─────────────────────────────────────\n",
    "if USE_API:\n",
    "    from tqdm.notebook import tqdm\n",
    "    import time\n",
    "\n",
    "    api_rows = []\n",
    "    RATE_LIMIT_DELAY = 1.5  # seconds between calls (free tier: 15 RPM)\n",
    "\n",
    "    for idx, row in tqdm(df.iterrows(), total=len(df), desc='Solving problems'):\n",
    "        t0 = time.time()\n",
    "        n_items = count_items(row['query'])\n",
    "        user_text = build_user_prompt(row['context'], row['query'], row.get('task_type', ''))\n",
    "\n",
    "        try:\n",
    "            response = gemini_model.generate_content(user_text)\n",
    "            raw = response.text.strip()\n",
    "        except Exception as e:\n",
    "            print(f'  ⚠ Error on row {idx}: {e}')\n",
    "            raw = ''\n",
    "\n",
    "            answers = parse_answers(raw, n_items)\n",
    "            explanation = extract_explanation(raw)\n",
    "\n",
    "        api_rows.append({\n",
    "            'id': row['id'],\n",
    "            'pred': json.dumps(answers, ensure_ascii=False),\n",
    "            'explanation': explanation,\n",
    "        })\n",
    "\n",
    "        elapsed = time.time() - t0\n",
    "        sleep_time = max(0, RATE_LIMIT_DELAY - elapsed)\n",
    "        if sleep_time > 0:\n",
    "            time.sleep(sleep_time)\n",
    "\n",
    "    api_submission = pd.DataFrame(api_rows)\n",
    "    api_submission.to_csv('submission_api.csv', index=False)\n",
    "    print(f'\\n✓ Wrote {len(api_rows)} rows → submission_api.csv')\n",
    "    api_submission.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3 — Local Model Mode (Qwen2.5-7B-Instruct, 4-bit)\n",
    "\n",
    "This is what runs in the **HuggingFace eval sandbox** (no internet, T4 GPU).  \n",
    "Enable a GPU runtime in Colab: Runtime → Change runtime type → T4 GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA available: True\n",
      "GPU: Tesla T4\n",
      "VRAM: 15.6 GB\n"
     ]
    }
   ],
   "source": [
    "# ── Check GPU availability ─────────────────────────────────────────────────\n",
    "import torch\n",
    "print(f'CUDA available: {torch.cuda.is_available()}')\n",
    "if torch.cuda.is_available():\n",
    "    print(f'GPU: {torch.cuda.get_device_name(0)}')\n",
    "    print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ Transformers stack ready\n"
     ]
    }
   ],
   "source": [
    "# ── Install model dependencies (skip if already installed) ─────────────────\n",
    "!pip install -q transformers accelerate autoawq\n",
    "print('✓ Transformers stack ready')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Will use model: Qwen/Qwen2.5-7B-Instruct\n",
      "Note: First download takes ~15 GB of disk space.\n"
     ]
    }
   ],
   "source": [
    "# ── Download Qwen2.5-14B-AWQ-Instruct (for local Colab testing) ────────────────\n",
    "# In the HF submission, the model is already in the repo root ('.')\n",
    "# Here we download it for local testing.\n",
    "\n",
    "LOCAL_MODEL_NAME = 'Qwen/Qwen2.5-14B-Instruct-AWQ'\n",
    "# Smaller alternative for RAM-limited environments:\n",
    "# LOCAL_MODEL_NAME = 'Qwen/Qwen2.5-3B-Instruct'\n",
    "\n",
    "print(f'Will use model: {LOCAL_MODEL_NAME}')\n",
    "print('Note: First download takes ~9 GB of disk space.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading tokeniser…\n",
      "Loading model in 4-bit (this may take 3–5 minutes)…\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dda4b131a8ef4555829dfd13ee92ad87",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/339 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ Model loaded!\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "\n",
    "print('Loading tokeniser…')\n",
    "tok = AutoTokenizer.from_pretrained(LOCAL_MODEL_NAME, trust_remote_code=True)\n",
    "\n",
    "print('Loading AWQ model…')\n",
    "local_model = AutoModelForCausalLM.from_pretrained(\n",
    "    LOCAL_MODEL_NAME,\n",
    "    torch_dtype=torch.float16,\n",
    "    device_map='auto',\n",
    "    trust_remote_code=True,\n",
    ").eval()\n",
    "\n",
    "print('✓ Model loaded!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In the sentences you provided, \"ka\" is a verb auxiliary that indicates the first person singular subject. Let's break down the sentences:\n",
      "\n",
      "1. **ŋa ka kɤ ne** = \"Do I go?\"\n",
      "2. **nɤ ka ku ne** = \"Does/Will you go?\"\n",
      "\n",
      "### Analysis:\n",
      "- **ŋa ka kɤ ne**: \n",
      "  - **ŋa** is likely a subject marker or a form of the verb \"to be\" in this context.\n",
      " \n"
     ]
    }
   ],
   "source": [
    "# ── Quick sanity test ─────────────────────────────────────────────────────\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "test_messages = [\n",
    "    {'role': 'system', 'content': 'You are an expert linguist.'},\n",
    "    {'role': 'user', 'content': 'What does \"ka\" mean if ŋa ka kɤ ne = \"Do I go?\" and nɤ ka ku ne = \"Does/Will you go?\"'}\n",
    "]\n",
    "\n",
    "test_inputs = tok.apply_chat_template(\n",
    "    test_messages, add_generation_prompt=True, return_tensors='pt', return_dict=True\n",
    ")\n",
    "test_inputs = {k: v.to(device) for k, v in test_inputs.items()}\n",
    "\n",
    "with torch.no_grad():\n",
    "    out = local_model.generate(**test_inputs, max_new_tokens=100, do_sample=False)\n",
    "\n",
    "print(tok.decode(out[0][test_inputs['input_ids'].shape[-1]:], skip_special_tokens=True))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d1131487f8047afb1fcb61cf881d62e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Local model inference:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  [1/3] id=012023020100 items=3 time=37.3s\n",
      "  [2/3] id=012023030200 items=3 time=12.4s\n",
      "  [3/3] id=012024010300 items=3 time=9.7s\n",
      "\n",
      "✓ Wrote 3 rows → submission.csv\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "  <div id=\"df-df0f69c3-cf30-46b1-80f0-de909dddfd62\" class=\"colab-df-container\">\n",
       "    <div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>pred</th>\n",
       "      <th>explanation</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>012023020100</td>\n",
       "      <td>[\"Will you sleep? / Do you sleep?\", \"Did I eat...</td>\n",
       "      <td>The sentences provided follow a pattern where ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>012023030200</td>\n",
       "      <td>[\"6\", \"16\", \"228\"]</td>\n",
       "      <td>From the given translations and the base 6 num...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>012024010300</td>\n",
       "      <td>[\"1. kite = fish → fishes = kitená\", \"2. nãka ...</td>\n",
       "      <td>From the examples, we can see that the plural ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "    <div class=\"colab-df-buttons\">\n",
       "      \n",
       "  <div class=\"colab-df-container\">\n",
       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-df0f69c3-cf30-46b1-80f0-de909dddfd62')\"\n",
       "            title=\"Convert this dataframe to an interactive table.\"\n",
       "            style=\"display:none;\">\n",
       "      \n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
       "  </svg>\n",
       "    </button>\n",
       "    \n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    .colab-df-buttons div {\n",
       "      margin-bottom: 4px;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "    <script>\n",
       "      const buttonEl =\n",
       "        document.querySelector('#df-df0f69c3-cf30-46b1-80f0-de909dddfd62 button.colab-df-convert');\n",
       "      buttonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "      async function convertToInteractive(key) {\n",
       "        const element = document.querySelector('#df-df0f69c3-cf30-46b1-80f0-de909dddfd62');\n",
       "        const dataTable =\n",
       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                    [key], {});\n",
       "        if (!dataTable) return;\n",
       "\n",
       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "          + ' to learn more about interactive tables.';\n",
       "        element.innerHTML = '';\n",
       "        dataTable['output_type'] = 'display_data';\n",
       "        await google.colab.output.renderOutput(dataTable, element);\n",
       "        const docLink = document.createElement('div');\n",
       "        docLink.innerHTML = docLinkHtml;\n",
       "        element.appendChild(docLink);\n",
       "      }\n",
       "    </script>\n",
       "  </div>\n",
       "  \n",
       "    </div>\n",
       "  </div>\n",
       "  "
      ],
      "text/plain": [
       "             id                                               pred  \\\n",
       "0  012023020100  [\"Will you sleep? / Do you sleep?\", \"Did I eat...   \n",
       "1  012023030200                                 [\"6\", \"16\", \"228\"]   \n",
       "2  012024010300  [\"1. kite = fish → fishes = kitená\", \"2. nãka ...   \n",
       "\n",
       "                                         explanation  \n",
       "0  The sentences provided follow a pattern where ...  \n",
       "1  From the given translations and the base 6 num...  \n",
       "2  From the examples, we can see that the plural ...  "
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ── Run local model inference ──────────────────────────────────────────────\n",
    "from tqdm.notebook import tqdm\n",
    "import time\n",
    "import json\n",
    "import torch\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "MAX_NEW_TOKENS = 1024\n",
    "local_rows = []\n",
    "\n",
    "for idx, row in tqdm(df.iterrows(), total=len(df), desc='Local model inference'):\n",
    "    t0 = time.time()\n",
    "    n_items = count_items(row['query'])\n",
    "    user_text = build_user_prompt(row['context'], row['query'], row.get('task_type', ''))\n",
    "\n",
    "    messages = [\n",
    "        {'role': 'system', 'content': SYSTEM_PROMPT},\n",
    "        {'role': 'user', 'content': user_text},\n",
    "    ]\n",
    "\n",
    "    model_inputs = tok.apply_chat_template(\n",
    "        messages, add_generation_prompt=True, return_tensors='pt', return_dict=True\n",
    "    )\n",
    "    model_inputs = {k: v.to(device) for k, v in model_inputs.items()}\n",
    "    input_ids = model_inputs['input_ids']\n",
    "\n",
    "    with torch.no_grad():\n",
    "        output_ids = local_model.generate(\n",
    "            **model_inputs,\n",
    "            max_new_tokens=MAX_NEW_TOKENS,\n",
    "            do_sample=False,\n",
    "            repetition_penalty=1.1,\n",
    "        )\n",
    "\n",
    "    raw = tok.decode(\n",
    "        output_ids[0][input_ids.shape[-1]:],\n",
    "        skip_special_tokens=True\n",
    "    ).strip()\n",
    "\n",
    "    answers = parse_answers(raw, n_items)\n",
    "    explanation = extract_explanation(raw)\n",
    "\n",
    "    local_rows.append({\n",
    "        'id': row['id'],\n",
    "        'pred': json.dumps(answers, ensure_ascii=False),\n",
    "        'explanation': explanation,\n",
    "    })\n",
    "\n",
    "    elapsed = time.time() - t0\n",
    "    print(f'  [{len(local_rows)}/{len(df)}] id={row[\"id\"]} items={n_items} time={elapsed:.1f}s')\n",
    "\n",
    "local_submission = pd.DataFrame(local_rows)\n",
    "local_submission.to_csv('submission.csv', index=False)\n",
    "print(f'\\n✓ Wrote {len(local_rows)} rows → submission.csv')\n",
    "local_submission.head()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 4 — Score Evaluation (Local)\n",
    "\n",
    "Estimate score on the mock data to iterate on your prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluating submission.csv…\n",
      "\n",
      "  [012023020100] item 1: pred=\"Will you sleep? / Do you sleep?\" ref=\"Does/Will you sleep?\" EM=0 chrF=0.670\n",
      "  [012023020100] item 2: pred=\"Did I eat?\" ref=\"Did I eat?\" EM=1 chrF=1.000\n",
      "  [012023020100] item 3: pred=\"Did you go?\" ref=\"Do you go?\" EM=0 chrF=0.593\n",
      "  [012023030200] item 1: pred=\"6\" ref=\"7\" EM=0 chrF=0.000\n",
      "  [012023030200] item 2: pred=\"16\" ref=\"20\" EM=0 chrF=0.000\n",
      "  [012023030200] item 3: pred=\"228\" ref=\"44\" EM=0 chrF=0.000\n",
      "  [012024010300] item 1: pred=\"1. kite = fish → fishes = kitená\" ref=\"kitea\" EM=0 chrF=0.259\n",
      "  [012024010300] item 2: pred=\"2. nãka = tree → trees = nãkaná\" ref=\"nãkana\" EM=0 chrF=0.334\n",
      "  [012024010300] item 3: pred=\"3. ãkitia = monkeys → monkey = ãkiti\" ref=\"ãkiti\" EM=0 chrF=0.379\n",
      "\n",
      "══════════════════════════════\n",
      "  Exact Match  : 0.1111\n",
      "  chrF         : 0.3594\n",
      "  Final Score  : 0.1998\n",
      "══════════════════════════════\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "import math\n",
    "import pandas as pd\n",
    "import sacrebleu\n",
    "\n",
    "# Fallback definition if Section 1 was not run\n",
    "if 'USE_MOCK' not in globals():\n",
    "    USE_MOCK = not os.path.exists('/tmp/data/test.csv')\n",
    "\n",
    "# Mock ground-truth answers (matching our mock problems)\n",
    "MOCK_ANSWERS = {\n",
    "    '012023020100': ['Does/Will you sleep?', 'Did I eat?', 'Do you go?'],\n",
    "    '012023030200': ['7', '20', '44'],\n",
    "    '012024010300': ['kitea', 'nãkana', 'ãkiti'],\n",
    "}\n",
    "\n",
    "\n",
    "def chrf_score(hypothesis: str, reference: str) -> float:\n",
    "    \"\"\"Compute chrF for a single hypothesis-reference pair.\"\"\"\n",
    "    result = sacrebleu.corpus_chrf([hypothesis], [[reference]])\n",
    "    return result.score / 100.0\n",
    "\n",
    "\n",
    "def evaluate_submission(submission_path: str, answers: dict) -> dict:\n",
    "    sub = pd.read_csv(submission_path, dtype=str).fillna('')\n",
    "    total_em = 0\n",
    "    total_chrf = 0\n",
    "    total_items = 0\n",
    "\n",
    "    for _, row in sub.iterrows():\n",
    "        pid = row['id']\n",
    "        if pid not in answers:\n",
    "            continue\n",
    "        refs = answers[pid]\n",
    "        \n",
    "        try:\n",
    "            preds = json.loads(row['pred'])\n",
    "        except Exception:\n",
    "            preds = []\n",
    "\n",
    "        # Pad preds to match refs length to ensure we score all items\n",
    "        while len(preds) < len(refs):\n",
    "            preds.append('')\n",
    "\n",
    "        for i, (pred, ref) in enumerate(zip(preds, refs)):\n",
    "            em = int(pred.strip().lower() == ref.strip().lower())\n",
    "            cf = chrf_score(pred, ref)\n",
    "            total_em += em\n",
    "            total_chrf += cf\n",
    "            total_items += 1\n",
    "            print(f'  [{pid}] item {i+1}: pred=\"{pred}\" ref=\"{ref}\" EM={em} chrF={cf:.3f}')\n",
    "\n",
    "    if total_items == 0:\n",
    "        return {'em': 0, 'chrf': 0, 'score': 0}\n",
    "\n",
    "    em_avg = total_em / total_items\n",
    "    chrf_avg = total_chrf / total_items\n",
    "    score = math.sqrt(em_avg * chrf_avg)\n",
    "    return {'em': em_avg, 'chrf': chrf_avg, 'score': score, 'n_items': total_items}\n",
    "\n",
    "\n",
    "if USE_MOCK:\n",
    "    # Evaluate whichever submission exists\n",
    "    sub_path = 'submission.csv' if os.path.exists('submission.csv') else 'submission_api.csv'\n",
    "    if os.path.exists(sub_path):\n",
    "        print(f'Evaluating {sub_path}…\\n')\n",
    "        results = evaluate_submission(sub_path, MOCK_ANSWERS)\n",
    "        print(f'\\n══════════════════════════════')\n",
    "        print(f'  Exact Match  : {results[\"em\"]:.4f}')\n",
    "        print(f'  chrF         : {results[\"chrf\"]:.4f}')\n",
    "        print(f'  Final Score  : {results[\"score\"]:.4f}')\n",
    "        print(f'══════════════════════════════')\n",
    "    else:\n",
    "        print('No submission file found. Run Section 2 or 3 first.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 5 — Inspect Outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ID: 12023020100\n",
      "  1. Will you sleep? / Do you sleep?\n",
      "  2. Did I eat?\n",
      "  3. Did you go?\n",
      "  [Explanation]: The sentences provided follow a pattern where the subject is indicated by the first word, followed by the verb, then a particle indicating past tense or future/will, and finally \"ne\" which seems to be\n",
      "\n",
      "ID: 12023030200\n",
      "  1. 6\n",
      "  2. 16\n",
      "  3. 228\n",
      "  [Explanation]: From the given translations and the base 6 number system, we can map the words directly to their numerical values.\n",
      "\n",
      "1. nif abo thonith = 6 * 1 = 6\n",
      "2. mer nif abo an = 2 * 6 + 4 = 12 + 4 = 16\n",
      "3. ithin \n",
      "\n",
      "ID: 12024010300\n",
      "  1. 1. kite = fish → fishes = kitená\n",
      "  2. 2. nãka = tree → trees = nãkaná\n",
      "  3. 3. ãkitia = monkeys → monkey = ãkiti\n",
      "  [Explanation]: From the examples, we can see that the plural form is created by adding \"-na\" to the stem of the word. For singular forms, it seems there might be a different pattern or no additional suffix for singu\n",
      "\n"
     ]
    }
   ],
   "source": [
    "sub_path = 'submission.csv' if os.path.exists('submission.csv') else 'submission_api.csv'\n",
    "if os.path.exists(sub_path):\n",
    "    sub = pd.read_csv(sub_path)\n",
    "    for _, row in sub.iterrows():\n",
    "        print(f'ID: {row[\"id\"]}')\n",
    "        preds = json.loads(row['pred'])\n",
    "        for i, p in enumerate(preds, 1):\n",
    "            print(f'  {i}. {p}')\n",
    "        print(f'  [Explanation]: {str(row.get(\"explanation\", \"\"))[:200]}')\n",
    "        print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Sync 14B AWQ model weights directly from Colab to Hugging Face ──────────\n",
    "# Run this cell inside Google Colab to download and sync the weights.\n",
    "# Colab's high network bandwidth completes this download + upload in seconds.\n",
    "\n",
    "from huggingface_hub import snapshot_download, HfApi\n",
    "\n",
    "import getpass\n",
    "TOKEN = getpass.getpass('Enter your Hugging Face WRITE token: ')\n",
    "SOURCE_REPO = 'Qwen/Qwen2.5-14B-Instruct-AWQ'\n",
    "TARGET_REPO = 'ahsanatiq98/iol-ai-submission'\n",
    "\n",
    "print('Downloading 14B AWQ model weights to Colab server cache...')\n",
    "local_path = snapshot_download(\n",
    "    repo_id=SOURCE_REPO,\n",
    "    token=TOKEN,\n",
    "    ignore_patterns=['*.git*', '*.gitattributes']\n",
    ")\n",
    "print('Download completed.')\n",
    "\n",
    "print(f'Uploading 14B AWQ model weights to target repo: {TARGET_REPO}...')\n",
    "api = HfApi(token=TOKEN)\n",
    "api.upload_folder(\n",
    "    folder_path=local_path,\n",
    "    repo_id=TARGET_REPO,\n",
    "    repo_type='model',\n",
    "    commit_message='Upload Qwen2.5-14B-Instruct-AWQ weights from Colab'\n",
    ")\n",
    "print('✓ Upload completed! Model weights are successfully synced.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 6 — HuggingFace Submission Prep\n",
    "\n",
    "Steps to submit to the competition:\n",
    "\n",
    "1. Create a **public** HuggingFace model repo\n",
    "2. Upload `script.py` (the standalone script, not this notebook)\n",
    "3. Download Qwen2.5-7B-Instruct weights into the repo:\n",
    "   ```bash\n",
    "   # In your HF repo, clone and run:\n",
    "   huggingface-cli download Qwen/Qwen2.5-7B-Instruct --local-dir . --local-dir-use-symlinks False\n",
    "   git add . && git commit -m \"Add model weights\" && git push\n",
    "   ```\n",
    "4. Submit your repo ID at the competition Space\n",
    "\n",
    "### Download the submission.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the submission file (in Google Colab)\n",
    "try:\n",
    "    from google.colab import files\n",
    "    if os.path.exists('submission.csv'):\n",
    "        files.download('submission.csv')\n",
    "        print('✓ Downloading submission.csv')\n",
    "    elif os.path.exists('submission_api.csv'):\n",
    "        files.download('submission_api.csv')\n",
    "        print('✓ Downloading submission_api.csv')\n",
    "    else:\n",
    "        print('No submission file found.')\n",
    "except ImportError:\n",
    "    print('Not in Colab. Find submission.csv in the current directory.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 📝 Notes & Tips\n",
    "\n",
    "### Improving Score\n",
    "- **Few-shot examples**: Add solved examples from previous IOL competitions to the system prompt\n",
    "- **Self-consistency**: Run the model N times and take majority vote\n",
    "- **Task-specific prompts**: Different prompt templates per `task_type`\n",
    "- **Post-processing**: Strip numbering from answers, normalise unicode\n",
    "\n",
    "### Rate Limits (Gemini Free Tier)\n",
    "- 15 requests/minute (RPM)\n",
    "- 1,500 requests/day (RPD)\n",
    "- Use `time.sleep(4)` between calls to stay safe\n",
    "\n",
    "### HuggingFace Sandbox Constraints\n",
    "- No internet at eval time\n",
    "- T4 GPU, 16 GB VRAM\n",
    "- 30-minute time limit\n",
    "- Pre-installed: `bitsandbytes`, `autoawq`, `transformers`, `pandas`, `torch`\n",
    "\n",
    "### Scoring Reminder\n",
    "$$\\text{score} = \\sqrt{\\text{EM}_w \\times \\text{chrF}_w}$$\n",
    "Both need to be high — exact matches matter, but partial credit counts too."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "name": "IOL-AI-2026-Solution.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}