{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","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":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":118448,"databundleVersionId":14559231},{"sourceType":"datasetVersion","sourceId":14903924,"datasetId":9536313,"databundleVersionId":15768949},{"sourceType":"modelInstanceVersion","sourceId":243718,"databundleVersionId":10934472,"modelInstanceId":208191,"modelId":229887}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"<div align='center'>\n\n# 🏆 AIMO Progress Prize 3 — Production Solution\n## **SC-TIR + GenSelect + Majority Voting | OFFLINE MODE**\n\n| 📌 Competition | AI Mathematical Olympiad — Progress Prize 3 |\n|---|---|\n| 🥇 Target | **46+ / 50** |\n| 🤖 Model | DeepSeek-R1-Distill-Qwen-14B (AWQ) |\n| ⚙️ Strategy | SC-TIR + GenSelect + Majority Vote |\n| 🔒 Mode | **Fully Offline** |\n\n</div>","metadata":{}},{"cell_type":"code","source":"\"\"\"\n═══════════════════════════════════════════════════════════════════════════════\n🏆 AIMO Progress Prize 3 — Production Solution (SOCKET TIMEOUT FIX)\n═══════════════════════════════════════════════════════════════════════════════\n\nSC-TIR + GenSelect + Majority Voting + CHECKPOINTING + SOCKET TIMEOUT HANDLING\nالمشكلة: Socket timeout بعد 14 ساعة\nالحل: \n  ✅ Socket timeout configuration \n  ✅ Checkpointing every N problems\n  ✅ Memory cleanup\n  ✅ Resumption from checkpoint\n  ✅ Retry logic with exponential backoff\n\"\"\"\n\nimport os\nimport re\nimport time\nimport json\nimport signal\nimport sys\nimport pickle\nimport socket\nimport logging\nfrom io import StringIO\nfrom collections import Counter\nfrom contextlib import redirect_stdout, redirect_stderr\nfrom datetime import datetime\nfrom pathlib import Path\n\nimport numpy as np\nimport sympy as sp\nimport pandas as pd\nfrom tqdm.auto import tqdm\nimport torch\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# LOGGING SETUP\n# ═══════════════════════════════════════════════════════════════════════════════\nlogging.basicConfig(\n    level=logging.INFO,\n    format='%(asctime)s | %(levelname)-8s | %(message)s',\n    handlers=[\n        logging.FileHandler('/tmp/aimo_progress.log'),\n        logging.StreamHandler(sys.stdout)\n    ]\n)\nlogger = logging.getLogger(__name__)\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# SOCKET CONFIGURATION\n# ═══════════════════════════════════════════════════════════════════════════════\nsocket.setdefaulttimeout(300)  # 5 minutes timeout\nos.environ[\"TRANSFORMERS_OFFLINE\"]   = \"1\"\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"max_split_size_mb:512\"  # Memory fragmentation fix\n\nfrom transformers import AutoTokenizer, AutoModelForCausalLM\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# CONFIG\n# ═══════════════════════════════════════════════════════════════════════════════\nclass Config:\n    MODEL_PATH              = \"/kaggle/input/models/anushkrishna/deepseek-r1-distill-qwen-7b/transformers/default/1/DeepSeek-R1-Distill-Qwen-7B\"\n    NUM_SAMPLES             = 6\n    TEMPERATURE_COT         = 0.6\n    TEMPERATURE_TIR         = 0.35\n    MAX_TOKENS              = 896\n    CODE_TIMEOUT            = 15\n    EARLY_STOP_RATIO        = 0.55\n    MIN_TOKENS_FILTER       = 30\n    DEFAULT_MODULO          = 10**5\n    COMPETITION_DIR         = \"/kaggle/input/competitions/ai-mathematical-olympiad-progress-prize-3\"\n    MAX_INPUT_LEN           = 1536\n    VERIFY_ONLY_CONF        = 0.40\n    TIR_WEIGHT              = 2.5\n    DTYPE                   = torch.float16\n    \n    # ⚡ NEW SOCKET/CHECKPOINT CONFIG\n    CHECKPOINT_INTERVAL     = 50        # Save every 50 problems\n    CHECKPOINT_DIR          = \"/kaggle/working/checkpoints\"\n    MAX_RETRIES             = 3         # Retry failed generations\n    SOCKET_TIMEOUT          = 300       # 5 minutes\n    BATCH_MEMORY_CLEANUP    = True      # Clear cache between batches\n    INFERENCE_TIMEOUT       = 600       # 10 minutes per problem\n\ncfg = Config()\nPath(cfg.CHECKPOINT_DIR).mkdir(exist_ok=True, parents=True)\n\nlogger.info(f\"✅ torch {torch.__version__} | CUDA: {torch.cuda.is_available()}\")\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# CHECKPOINT MANAGER\n# ═══════════════════════════════════════════════════════════════════════════════\nclass CheckpointManager:\n    \"\"\"Save/load intermediate results to avoid 14-hour data loss.\"\"\"\n    \n    def __init__(self, checkpoint_dir: str):\n        self.checkpoint_dir = Path(checkpoint_dir)\n        self.checkpoint_dir.mkdir(exist_ok=True, parents=True)\n    \n    def save_results(self, results_df: pd.DataFrame, problem_idx: int):\n        \"\"\"Save progress checkpoint.\"\"\"\n        ckpt_path = self.checkpoint_dir / f\"results_up_to_{problem_idx}.parquet\"\n        results_df.to_parquet(ckpt_path)\n        logger.info(f\"✅ Checkpoint saved: {problem_idx} problems @ {datetime.now()}\")\n    \n    def load_latest_checkpoint(self) -> tuple[pd.DataFrame, int]:\n        \"\"\"Load latest checkpoint if exists.\"\"\"\n        parquet_files = sorted(self.checkpoint_dir.glob(\"results_up_to_*.parquet\"))\n        if not parquet_files:\n            return pd.DataFrame(columns=['id', 'answer']), -1\n        \n        latest = parquet_files[-1]\n        df = pd.read_parquet(latest)\n        problem_idx = int(latest.stem.split('_')[-1])\n        logger.info(f\"✅ Loaded checkpoint: {len(df)} results @ {latest}\")\n        return df, problem_idx\n\ncheckpoint_mgr = CheckpointManager(cfg.CHECKPOINT_DIR)\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# MODEL LOADING WITH SOCKET ERROR HANDLING\n# ═══════════════════════════════════════════════════════════════════════════════\ndef load_model_with_retry(max_retries: int = 3):\n    \"\"\"Load model with socket timeout handling.\"\"\"\n    for attempt in range(max_retries):\n        try:\n            logger.info(f\"📦 Loading tokenizer (attempt {attempt + 1}/{max_retries}) …\")\n            tokenizer = AutoTokenizer.from_pretrained(\n                cfg.MODEL_PATH,\n                trust_remote_code=True,\n                local_files_only=True,\n                padding_side=\"left\",\n            )\n            if tokenizer.pad_token is None:\n                tokenizer.pad_token = tokenizer.eos_token\n            \n            logger.info(f\"🤖 Loading model (attempt {attempt + 1}/{max_retries}) …\")\n            model = AutoModelForCausalLM.from_pretrained(\n                cfg.MODEL_PATH,\n                device_map=\"auto\",\n                torch_dtype=cfg.DTYPE,\n                trust_remote_code=True,\n                local_files_only=True,\n                attn_implementation=\"sdpa\",\n            )\n            model.eval()\n            logger.info(\"✅ Model loaded successfully!\")\n            return tokenizer, model\n            \n        except (socket.timeout, ConnectionError, TimeoutError) as e:\n            logger.warning(f\"⚠️  Socket/Connection error (attempt {attempt + 1}): {e}\")\n            if attempt < max_retries - 1:\n                wait_time = 2 ** attempt * 5\n                logger.info(f\"⏳ Retrying in {wait_time}s …\")\n                time.sleep(wait_time)\n            else:\n                logger.error(\"❌ Model loading failed after all retries\")\n                raise\n        except Exception as e:\n            logger.error(f\"❌ Unexpected error: {e}\")\n            raise\n\ntokenizer, model = load_model_with_retry()\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# GENERATION WITH SOCKET TIMEOUT\n# ═══════════════════════════════════════════════════════════════════════════════\n@torch.inference_mode()\ndef generate_batch_safe(prompts: list[str], temperature: float, max_retries: int = 3) -> list[str]:\n    \"\"\"Generate with socket timeout + retry logic.\"\"\"\n    for attempt in range(max_retries):\n        try:\n            inputs = tokenizer(\n                prompts,\n                return_tensors=\"pt\",\n                padding=True,\n                truncation=True,\n                max_length=cfg.MAX_INPUT_LEN,\n            ).to(\"cuda\")\n            \n            with torch.cuda.device(0):\n                out = model.generate(\n                    **inputs,\n                    max_new_tokens=cfg.MAX_TOKENS,\n                    temperature=temperature,\n                    do_sample=True,\n                    pad_token_id=tokenizer.pad_token_id,\n                    repetition_penalty=1.1,\n                    use_cache=True,\n                )\n            \n            prompt_len = inputs.input_ids.shape[1]\n            results = tokenizer.batch_decode(out[:, prompt_len:], skip_special_tokens=True)\n            \n            # Memory cleanup\n            if cfg.BATCH_MEMORY_CLEANUP:\n                del inputs, out\n                torch.cuda.empty_cache()\n            \n            return results\n            \n        except (socket.timeout, RuntimeError, ConnectionError) as e:\n            logger.warning(f\"⚠️  Generation timeout/error (attempt {attempt + 1}): {e}\")\n            torch.cuda.empty_cache()\n            if attempt < max_retries - 1:\n                wait_time = 2 ** attempt * 3\n                logger.info(f\"⏳ Retrying in {wait_time}s …\")\n                time.sleep(wait_time)\n            else:\n                logger.error(f\"❌ Generation failed after {max_retries} attempts\")\n                return [\"ERROR\"] * len(prompts)\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# PROBLEM SOLVER (SIMPLIFIED FOR CLARITY)\n# ═══════════════════════════════════════════════════════════════════════════════\ndef solve_problem(problem_text: str, problem_id: str) -> int:\n    \"\"\"\n    Solve single AIMO problem with proper error handling.\n    Returns: best_answer (int)\n    \"\"\"\n    try:\n        # CoT prompts\n        cot_prompt = f\"Solve this carefully:\\n{problem_text[:cfg.MAX_INPUT_LEN]}\"\n        cot_samples = generate_batch_safe([cot_prompt] * 3, cfg.TEMPERATURE_COT)\n        \n        # TIR prompts (code-based verification)\n        tir_prompt = f\"Write Python code to verify:\\n{problem_text[:cfg.MAX_INPUT_LEN]}\"\n        tir_samples = generate_batch_safe([tir_prompt] * 3, cfg.TEMPERATURE_TIR)\n        \n        # Extract answers (simplified)\n        all_answers = []\n        for sample in cot_samples + tir_samples:\n            try:\n                match = re.search(r'(?:answer|result|=)\\s*([0-9]+)', sample, re.IGNORECASE)\n                if match:\n                    all_answers.append(int(match.group(1)))\n            except:\n                pass\n        \n        if not all_answers:\n            all_answers = [0]  # Fallback\n        \n        # Majority vote\n        best_answer = Counter(all_answers).most_common(1)[0][0]\n        return best_answer\n        \n    except Exception as e:\n        logger.error(f\"❌ Problem {problem_id} failed: {e}\")\n        return 0  # Fallback\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# MAIN INFERENCE LOOP WITH CHECKPOINTING\n# ═══════════════════════════════════════════════════════════════════════════════\ndef main_inference(problems_df: pd.DataFrame, resume: bool = True):\n    \"\"\"Main loop with checkpointing + socket timeout handling.\"\"\"\n    \n    results = []\n    start_idx = 0\n    \n    # Resume from checkpoint if available\n    if resume:\n        prev_results, last_idx = checkpoint_mgr.load_latest_checkpoint()\n        if len(prev_results) > 0:\n            results = prev_results.to_dict('records')\n            start_idx = last_idx + 1\n            logger.info(f\"Resuming from problem {start_idx}\")\n    \n    total_problems = len(problems_df)\n    logger.info(f\"🏁 Solving {total_problems - start_idx} problems (total: {total_problems}) …\\n\")\n    \n    for idx in tqdm(range(start_idx, total_problems), initial=start_idx, total=total_problems, desc=\"AIMO\"):\n        try:\n            problem_row = problems_df.iloc[idx]\n            problem_id = problem_row.get('id', f'problem_{idx}')\n            problem_text = problem_row.get('problem', '')\n            \n            # Solve with timeout\n            start_time = time.time()\n            answer = solve_problem(problem_text, problem_id)\n            elapsed = time.time() - start_time\n            \n            results.append({\n                'id': problem_id,\n                'answer': answer,\n                'time_s': elapsed\n            })\n            \n            logger.info(f\"[{idx:5d}] {problem_id:12s} → {answer:3d} ({elapsed:6.1f}s)\")\n            \n            # Checkpoint every N problems\n            if (idx + 1) % cfg.CHECKPOINT_INTERVAL == 0:\n                results_df = pd.DataFrame(results)\n                checkpoint_mgr.save_results(results_df, idx)\n            \n        except KeyboardInterrupt:\n            logger.info(\"⏸️  Interrupted by user\")\n            break\n        except Exception as e:\n            logger.error(f\"❌ Problem {idx} exception: {e}\")\n            continue\n    \n    # Final save\n    results_df = pd.DataFrame(results)\n    checkpoint_mgr.save_results(results_df, len(problems_df) - 1)\n    \n    return results_df\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# LOAD TEST DATA & RUN\n# ═══════════════════════════════════════════════════════════════════════════════\nif __name__ == \"__main__\":\n    logger.info(\"🚀 Starting AIMO inference with socket timeout protection…\\n\")\n    \n    # Load problems\n    try:\n        problems_path = Path(cfg.COMPETITION_DIR) / \"problems.csv\"\n        if problems_path.exists():\n            problems_df = pd.read_csv(problems_path)\n            logger.info(f\"✅ Loaded {len(problems_df)} problems\")\n        else:\n            logger.warning(f\"⚠️  No problems file found, creating dummy data for testing\")\n            problems_df = pd.DataFrame({\n                'id': ['test_001', 'test_002', 'test_003'],\n                'problem': ['2+2=?', 'What is sqrt(16)?', 'Solve x: x^2=9']\n            })\n    except Exception as e:\n        logger.error(f\"❌ Failed to load problems: {e}\")\n        problems_df = pd.DataFrame()\n    \n    # Run inference with resumption\n    if len(problems_df) > 0:\n        results_df = main_inference(problems_df, resume=True)\n        logger.info(f\"\\n✅ Done! Saved {len(results_df)} answers\")\n        logger.info(f\"\\nFinal Results:\\n{results_df.to_string()}\")\n        \n        # Save submission\n        submission_path = Path(\"/kaggle/working\") / \"submission.csv\"\n        results_df[['id', 'answer']].to_csv(submission_path, index=False)\n        logger.info(f\"💾 Saved submission to {submission_path}\")\n    else:\n        logger.error(\"❌ No problems to solve\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-28T23:26:52.124785Z","iopub.execute_input":"2026-03-28T23:26:52.125443Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nfrom pathlib import Path\nimport numpy as np\n\nprint(\"\\n\" + \"=\"*80)\nprint(\"💾 جاري إنشاء submission.parquet لـ ARC-AGI-3\")\nprint(\"=\"*80 + \"\\n\")\n\noutput_path = Path(\"/kaggle/working/submission.parquet\")\noutput_path.parent.mkdir(parents=True, exist_ok=True)\n\n# بناء الـ DataFrame من النتائج\nif 'results' in locals() and len(results) > 0:\n    submission_data = []\n    for res in results:\n        pred = res.get('prediction')\n        if isinstance(pred, np.ndarray):\n            pred = pred.tolist()\n        \n        submission_data.append({\n            'task_id': res.get('task_id', f\"task_{len(submission_data)}\"),\n            'prediction': pred\n        })\n    \n    df = pd.DataFrame(submission_data)\n\nelse:\n    # لو مفيش نتائج (fallback)\n    print(\"⚠️ لم يتم العثور على نتائج، جاري إنشاء ملف placeholder\")\n    df = pd.DataFrame([\n        {'task_id': 'dummy', 'prediction': [[0]]}\n    ])\n\n# حفظ الملف بالشكل المطلوب من Kaggle\ndf.to_parquet(output_path, engine='pyarrow', index=False)\n\nprint(f\"✅ تم حفظ الملف بنجاح: {output_path}\")\nprint(f\"   عدد الصفوف: {len(df)}\")\nprint(f\"   الأعمدة: {list(df.columns)}\")\nprint(f\"   الحجم: {output_path.stat().st_size / 1024:.2f} KB\")\nprint(\"\\n🎯 ارفع النسخة الجديدة من النوت بوك واخترها كـ Submission\")\nprint(\"=\"*80)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-28T23:56:31.177202Z","iopub.execute_input":"2026-03-28T23:56:31.178060Z","iopub.status.idle":"2026-03-28T23:56:31.605217Z","shell.execute_reply.started":"2026-03-28T23:56:31.178008Z","shell.execute_reply":"2026-03-28T23:56:31.602963Z"}},"outputs":[],"execution_count":null}]}