Upload platform_agnostic_inference_code_score21.py

platform_agnostic_inference_code_score21.py

@@ -0,0 +1,912 @@
+# -*- coding: utf-8 -*-
+"""Final-Submission -gpt-oss-score21.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1hMIhGYkVJqyB_Qv_GLgH3d21hInIvKZp
+
+# Setup The Environment
+"""
+
+# Commented out IPython magic to ensure Python compatibility.
+# %%bash
+# pip install paramiko math_verify litellm flashinfer-python vllm==0.11.2 openai_harmony
+# 
+# pip install absl-py==2.4.0 \
+#     catalogue==2.0.10 \
+#     colorful==0.5.8 \
+#     contextlib2==21.6.0 \
+#     decorator==5.2.1 \
+#     deprecated==1.3.1 \
+#     distlib==0.4.0 \
+#     docker==7.1.0 \
+#     exceptiongroup==1.3.1 \
+#     fabric==3.2.2 \
+#     fiddle==0.3.0 \
+#     google-api-core==2.29.0 \
+#     google-auth==2.48.0 \
+#     googleapis-common-protos==1.72.0 \
+#     graphviz==0.21 \
+#     grpcio==1.78.0 \
+#     h2==4.3.0 \
+#     hf-xet==1.2.0 \
+#     hpack==4.1.0 \
+#     hyperframe==6.1.0 \
+#     inquirerpy==0.3.4 \
+#     ledoc-ui==0.1.0 \
+#     leptonai==0.27.0 \
+#     libcst==1.8.6 \
+#     mypy-extensions==1.1.0 \
+#     nemo-run==0.6.0 \
+#     omegaconf==2.3.0 \
+#     opencensus==0.11.4 \
+#     opencensus-context==0.1.3 \
+#     opentelemetry-api==1.39.1 \
+#     opentelemetry-exporter-prometheus==0.60b1 \
+#     opentelemetry-proto==1.39.1 \
+#     opentelemetry-sdk==1.39.1 \
+#     opentelemetry-semantic-conventions==0.60b1 \
+#     pfzy==0.3.4 \
+#     platformdirs==4.9.2 \
+#     prompt-toolkit==3.0.52 \
+#     proto-plus==1.27.1 \
+#     py-spy==0.4.1 \
+#     pyasn1==0.6.2 \
+#     pyasn1-modules==0.4.2 \
+#     pyre-extensions==0.0.32 \
+#     python-multipart==0.0.22 \
+#     rsa==4.9.1 \
+#     smart-open==7.5.0 \
+#     toml==0.10.2 \
+#     torchx==0.7.0 \
+#     typer-slim==0.24.0 \
+#     virtualenv==20.37.0 \
+#     wcwidth==0.6.0 \
+#     wrapt==2.1.1
+# 
+# pip install openpyxl
+#
+
+# Track Overall Time
+import time
+global_deadline = time.perf_counter() + 5*3600
+global_remaining = global_deadline - time.perf_counter()
+cutoff_duration = global_remaining - 350
+def get_global_remaining():
+    return max(0, global_deadline - time.perf_counter())
+
+import os
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+import torch
+
+import asyncio
+import torch
+import subprocess
+import warnings
+import glob
+import pandas as pd
+import traceback
+import nest_asyncio
+import httpx
+import re
+import time
+import copy
+import json
+import requests
+import pandas as pd
+import polars as pl
+from collections import Counter
+from typing import List
+import secrets
+import json
+pd.set_option('display.max_colwidth', None)
+warnings.filterwarnings("ignore", category=SyntaxWarning)
+nest_asyncio.apply()
+os.environ["TORCH_COMPILE_DISABLE"] = "1"
+os.environ["TORCHDYNAMO_DISABLE"] = "1"
+os.environ['TRANSFORMERS_NO_FLAX'] = '1'
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+os.environ['TRITON_PTXAS_PATH'] = '/usr/local/cuda/bin/ptxas'
+os.environ['TIKTOKEN_RS_CACHE_DIR']= "/content/harmony_encoding"
+os.environ["TORCH_CUDA_ARCH_LIST"] = '9.0'
+os.environ["PYTORCH_CUDA_ALLOC_CONF"]="expandable_segments:True"
+#os.environ["VLLM_USE_FLASHINFER_SAMPLER"]= "1"
+from collections import Counter, defaultdict
+
+# This will change in kaggle
+os.environ["TORCHINDUCTOR_CACHE_DIR"] = "torch_cache"
+
+import os, sys
+original_pythonpath = os.environ.get("PYTHONPATH", "")
+path1 = '/content/modified-nemo-skills'
+merged_pythonpath = f"{path1}:{original_pythonpath}" if original_pythonpath else {path1}
+os.environ["PYTHONPATH"] = merged_pythonpath
+sys.path.append('/content/modified-nemo-skills')
+
+from nemo_skills.code_execution.sandbox import get_sandbox
+from nemo_skills.inference.model import get_code_execution_model
+from nemo_skills.prompt.utils import get_prompt
+from nemo_skills.inference.model import get_model
+
+"""# Configuration Parameters"""
+
+host = "127.0.0.1"
+port = 5000
+tp_size = 1
+max_public = 10
+max_tokens = 38000
+max_input_tokens = 2050
+tokens_to_generate =  35950 - 10
+max_batch_size = 8
+timeout_seconds = 300
+global_buffer = 350
+finish_at_last_n = 2
+max_code_output_characters = 1100
+code_execution_timeout = 10
+max_code_executions = 125
+g_score = 0
+g_count = 0
+prompt_score = Counter()
+sampling_params = {
+    "tokens_to_generate": tokens_to_generate,
+    "temperature": 1, # 0.2,
+    "top_p": 1,
+}
+
+thoughts = [""] * 50
+thoughts = thoughts[:max_batch_size]
+i = 0
+
+model_path = "/content/model"
+
+"""# Start Server - Load Model & Sandbox"""
+
+server_started = False
+def load_model():
+    cmd = [
+        "python",
+        "-m",
+        "nemo_skills.inference.server.serve_vllm",
+        f"--model={model_path}",
+        "--port=5000",
+        "--num_gpus=1",
+        "--max-model-len=38000",
+        "--max-num-batched-tokens=16384",
+        "--max-num-seqs=11",
+        "--max-cudagraph-capture-size=2048",
+        "--gpu-memory-utilization=0.95",
+        "--kv-cache-dtype=auto",
+        "--stream-interval=200",
+        "--enable-prefix-caching",
+        "--uvicorn-log-level debug",
+        "--enable-log-requests",
+        "--enable-log-outputs",
+        "--async-scheduling",
+         ]
+
+    log_file = open("vllm.log", "w")
+    vllm_server = subprocess.Popen(
+      cmd,
+      stdout=log_file,
+      stderr=log_file,
+      text=True,
+      bufsize=1  # line-buffered
+    )
+    return vllm_server
+
+vllm_server=load_model()
+
+def wait_for_server(url=f"http://{host}:{port}", timeout=1400):
+    start = time.perf_counter()
+    while True:
+        try:
+            r = requests.get(f"{url}/docs")
+            if r.status_code == 200:
+                print("✅ Server is ready",time.perf_counter()-start)
+                return True
+        except Exception:
+            pass
+
+        if time.perf_counter() - start > timeout:
+            raise TimeoutError("Server did not start in time")
+
+        time.sleep(1)
+
+def sandbox_server():
+    log_file = open("sandbox.log", "w")
+    sandbox_process = subprocess.Popen(
+        ["python", "-m", "nemo_skills.code_execution.local_sandbox.local_sandbox_server"],
+                stdout=log_file,
+                stderr=log_file,
+                text=True,
+                bufsize=1)
+
+    time.sleep(3)
+
+time.sleep(2)
+sandbox_server()
+sandbox = get_sandbox()  # localhost by default
+
+"""# Prompt Types and Updating Prompt"""
+
+default_prompt = (
+        'You are an elite mathematical problem solver with expertise at the International '
+        'Mathematical Olympiad (IMO) level. Your goal is to find the correct answer through '
+        'rigorous mathematical reasoning.\n\n'
+
+        '# Problem-Solving Approach:\n'
+        '1. UNDERSTAND: Carefully read and rephrase the problem in your own words. '
+        'Identify what is given, what needs to be found, and any constraints.\n'
+        '2. EXPLORE: Consider multiple solution strategies. Think about relevant theorems, '
+        'techniques, patterns, or analogous problems. Don\'t commit to one approach immediately.\n'
+        '3. PLAN: Select the most promising approach and outline key steps before executing.\n'
+        '4. EXECUTE: Work through your solution methodically. Show all reasoning steps clearly.\n'
+        '5. VERIFY: Check your answer by substituting back, testing edge cases, or using '
+        'alternative methods. Ensure logical consistency throughout.\n\n'
+
+        '# Mathematical Reasoning Principles:\n'
+        '- Break complex problems into smaller, manageable sub-problems\n'
+        '- Look for patterns, symmetries, and special cases that provide insight\n'
+        '- Use concrete examples to build intuition before generalizing\n'
+        '- Consider extreme cases and boundary conditions\n'
+        '- If stuck, try working backwards from the desired result\n'
+        '- Be willing to restart with a different approach if needed\n\n'
+
+        '# Verification Requirements:\n'
+        '- Cross-check arithmetic and algebraic manipulations\n'
+        '- Verify that your solution satisfies all problem constraints\n'
+        '- Test your answer with simple cases or special values when possible\n'
+        '- Ensure dimensional consistency and reasonableness of the result\n\n'
+
+        "#RESPONSE FORMAT:\n\n"
+        "The final answer must be a non-negative integer.\n. Instead of the \\boxed{} format use json format. Follow the instructions for the format-"
+        ' "Answer": <non-negative integer>,"Confidence": <number between 0 and 1>'
+        "Do not output any additional reasoning after this JSON.\n"
+    )
+
+answerext_prompt = (
+            "You are an answer extraction system for competition mathematics problems "
+            "(olympiad, Putnam, HMMT, AMC/AIME style and similar).\n"
+            "You will be given a PROBLEM and a MODEL_RESPONSE. The response may be "
+            "incomplete, truncated mid-reasoning, or cut off before a formal conclusion.\n\n"
+            "YOUR JOB:\n"
+            "1. Read the problem and identify the SOUGHT QUANTITY — this could be a "
+            "numerical value, an expression, a set of solutions, a function, a "
+            "characterization, a bound, a count, an extremal quantity, a geometric "
+            "measure, or a closed-form answer that the problem asks to find, "
+            "determine, or compute.\n"
+            "2. Extract the model's best answer to that sought quantity from the "
+            "response, even if:\n"
+            "   - It is not explicitly labeled with 'the answer is' or 'therefore'\n"
+            "   - It appears mid-sentence or mid-calculation\n"
+            "   - The response was truncated before a formal conclusion\n"
+            "   - It is stated with hedging language like 'seems to be,' 'so we get,' "
+            "or 'this gives'\n\n"
+            "RULES FOR IDENTIFYING THE ANSWER:\n"
+            "- The model's reasoning often explores multiple cases, subcases, or "
+            "candidate values. Distinguish between:\n"
+            "    (a) INTERMEDIATE SUB-RESULTS: values computed within a single case "
+            "or step (e.g., 'LCM 60,' 'sum = 14,' 'this gives 42') that feed "
+            "into the broader argument but do not directly answer the problem.\n"
+            "    (b) THE CANDIDATE ANSWER: the value, expression, or characterization "
+            "of the sought quantity that the model is building toward or "
+            "accumulating evidence for across its reasoning.\n"
+            "  Extract (b), not (a).\n"
+            "- In optimization or extremal problems, the model may test many "
+            "configurations and compare them against a leading candidate value. "
+            "If the model repeatedly checks whether alternatives can 'beat,' "
+            "'exceed,' or 'improve upon' a particular value, and none do, "
+            "treat that value as the candidate answer — even if the response "
+            "ends before a formal conclusion.\n"
+            "- If the model arrives at the same value through multiple approaches "
+            "or repeatedly returns to it as the best result, that is strong "
+            "signal it is the intended answer.\n"
+            "- If the response is truncated but a candidate answer is visible "
+            "from the reasoning so far, extract it. A truncated response with "
+            "a clear leading candidate is better than no answer.\n"
+            "- If the problem asks to 'compute,' 'find,' or 'determine' a specific "
+            "quantity, look for the last/best concrete value of THAT SPECIFIC "
+            "quantity — not intermediate quantities used along the way.\n"
+            "- If the answer appears in LaTeX formatting (e.g., \\boxed{140}, "
+            "$\\frac{7}{3}$, or similar), extract the value inside the formatting.\n\n"
+            "RESPONSE FORMAT:\n"
+            "Your ONLY task is to output a single JSON object — no preamble, no explanation, no mathematical calculations.\n"
+            "The final answer must be a non-negative integer. Instead of the \\boxed{{}} format use json format. Follow the instructions for the format-"
+            ' {{"Answer": <non-negative integer>,"Confidence": <number between 0 and 1>}}'
+            "Do not output any additional reasoning after this JSON.\n"
+        )
+
+# Below will change
+system_message='{system_prompt}'
+prompt_template = get_prompt(prompt_config='gpt-oss/math',system_message=system_message,tokenizer=model_path,code_tags="gpt-oss")
+chat_template_kwargs = {
+    "builtin_tools": ["python"],
+    "reasoning_effort":"high"
+
+}
+
+def safe_concat(a, b,function_name):
+    if a is None or b is None:
+        raise ValueError(f"Cannot concatenate: a={a}, b={b}, Error Raised from function {function_name}")
+    return a + b
+
+"""# Data Extraction & Early Stopping"""
+
+class Result:
+    def __init__(self):
+        self.early_stop_flag = False
+    def best_voted_answer(self):
+        return self.best_answer
+
+    def majority_voting(self, answer_list):
+        count = defaultdict(float)
+        # Keep raw list separate; filter into valid_answers
+        self.answer_list = answer_list
+        self.valid_answers = [x["Answer"] for x in self.answer_list if x["Answer"] != -1]
+        print("Answer_list after popping -1", self.valid_answers, "%%%%")
+
+        # BUG FIX: set fallback when all answers are invalid
+        if len(self.valid_answers) == 0:
+            self.best_answer = None
+            self.best_count = 0
+            self.second_count = 0
+            self.sorted_answers = []
+            return
+
+        for a in self.valid_answers:
+            count[a] += 1
+        self.sorted_answers = sorted(count.items(), key=lambda x: x[1], reverse=True)
+
+        self.best_answer, self.best_count = self.sorted_answers[0]
+        self.second_count = self.sorted_answers[1][1] if len(self.sorted_answers) > 1 else 0
+
+        if (
+            self.best_count == 1
+            and self.best_answer == 0
+            and len(self.sorted_answers) > 1
+            and self.sorted_answers[1] is not None
+        ):
+
+            self.best_answer, self.best_count = self.sorted_answers[1]
+
+
+    def early_stop(self, answer_list, num_done):
+        print("Num_done is",num_done)
+        self.num_done = num_done
+        self.majority_voting(answer_list)
+        n_valid = len(self.valid_answers)
+        best = self.best_count
+        gap = self.best_count - self.second_count
+        print(f"Num done: {self.num_done}, Valid answers: {n_valid}, "
+              f"Best count: {best}, Second count: {self.second_count}")
+
+        if n_valid == 0:
+            return False
+
+        if best >= 4 and gap >= 2:
+            self.early_stop_flag = True
+            print(f">>> EARLY STOP at {self.num_done} completions | "
+                  f"best={self.best_answer} (count={best}, gap={gap})")
+
+        return self.early_stop_flag
+
+    def get_best_answer(self,answer_list, num_done, flag):
+        if not flag:
+            self.majority_voting(answer_list)
+        else:
+            self.early_stop(answer_list, num_done)
+        return self.best_voted_answer(), self.early_stop_flag
+
+import re, requests
+
+class Answer:
+    def __init__(self):
+        self.best_answer = None
+        self.input_message  = ""
+        self.best_count = 0
+        self.second_count = 0
+        self.answer_list = []          # ← was None, init as empty list
+        self.early_stop_flag = False
+        self.sorted_answers = []
+        self.valid_answers = []        # ← filtered list (no -1s), kept separate
+        self.sampling_param = {
+            "tokens_to_generate": 7000,
+            "temperature": 0.9, # 0.2,
+            "top_p": 0.95,
+             }
+        self.timeout = httpx.Timeout(
+           connect=60.0,
+           read=300.0,
+           write=60.0,
+           pool=120.0,
+        )
+
+    def clean_messages(self, text):
+        cleaned = re.sub(r'<\|[^|]*\|>', '', text)
+        return cleaned.strip()
+
+
+    async def extract_answer(self, question, model_output):
+        answer = -1
+        confidence = -0.1
+        seed = secrets.randbits(32)
+        input_message = self.clean_messages(model_output)
+        rid = secrets.token_hex(8)
+        message = prompt_template.fill(
+                input_dict={
+                    "problem": safe_concat(question,input_message,"extract_answer"),
+                    "system_prompt": answerext_prompt,
+                },
+                chat_template_kwargs = chat_template_kwargs,
+                format_as_string=True
+            )
+        print(prompt_template)
+        print("textd was called")
+        try:
+            data, completion_tokens = await server_obj.generate_response(
+                prompt=message,
+                random_seed=seed,
+                stream=True,
+                calling_function = "extract_answer",
+                extra_body={"request_id": rid},
+                timeout = self.timeout,
+               **self.sampling_param,
+             )
+
+            if data is not None and isinstance(data, dict):
+                return data
+            else:
+                return {"Answer":-1, "Confidence":-0.1}
+
+        except Exception as e:
+            print(f"[extract_answer failed] {type(e).__name__}: {e}")
+            return {"Answer":answer,"Confidence": confidence}
+
+"""# Inference"""
+
+# Below will change in kaggle
+#Instantiate Server Object
+server_obj = get_code_execution_model(server_type = 'vllm',
+                                model=model_path,
+                               base_url="http://127.0.0.1:5000/v1",
+                               api_key='EMPTY',
+                               sandbox=sandbox,
+                               code_execution={
+                                'max_code_output_characters': max_code_output_characters,
+                                'code_execution_timeout': code_execution_timeout,
+                                'max_code_executions': max_code_executions,
+                               })
+
+async def abort_request(request_ids: str | list[str]):
+    """Sequential best-effort server-side abort.
+    Uses short timeouts so a slow/down server doesn't block.
+    Silently ignores failures.
+    """
+    if isinstance(request_ids, str):
+        request_ids = [request_ids]
+
+    timeout = httpx.Timeout(connect=1.0, read=2.0, write=1.0, pool=1.0)
+
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        for rid in request_ids:
+            try:
+                await client.delete(f"http://{host}:{port}/v1/requests/{rid}")
+            except Exception:
+                # optionally log instead of silent pass
+                pass
+            await asyncio.sleep(0.05)  # cooperative yield
+
+class ClientClass:
+    def __init__(self, prompt):
+        global sampling_params
+        self.thresh_hold = 4                         # minimum completions before checking early stop
+        self.system_prompt = prompt
+        self.answer = {}
+        self.randomseed_list = []
+        self.num_done = 0
+        self.sampling_param = copy.deepcopy(sampling_params)
+        self.question = ""
+        self.finished_generations = []
+        self.final_answer = None
+        self.early_stop_flag = False
+        self.flattened_prompt_list = []
+        self.list_of_questions = []
+        self.answer_list = []
+        self.request_ids = []  # per-task IDs for server-side abort
+        self.tasks = []
+        self.timeout = httpx.Timeout(
+            connect=30.0,
+            read= 500.0 ,
+            write=30.0,
+            pool=120.0,
+        )
+        self.answerobj = Answer()
+
+    async def send_request_to_server(self):
+        print("Request sent")
+        self.request_ids = [secrets.token_hex(8) for _ in self.list_of_questions]
+        self.randomseed_list = [k for k in range(len(self.list_of_questions))]
+        for prompt, seed, rid in zip(self.list_of_questions, self.randomseed_list, self.request_ids):
+            task = asyncio.create_task(
+                      server_obj.generate_async(
+                        prompt=prompt,
+                        random_seed=seed,
+                        timeout=self.timeout,
+                        remove_stop_phrases=False,
+                        stream = True,
+                        extra_body={"request_id": rid},
+                        **prompt_template.get_code_execution_args(),
+                        **self.sampling_param,
+                         )
+                      )
+            self.tasks.append(task)
+
+        try:
+            processed = set()
+            for completed in asyncio.as_completed(self.tasks):
+                try:
+                    result = await completed
+                    print("Total number of generated tokens", result["total_num_generated_tokens"])
+                    self.num_done += 1
+                    processed.add(completed)   # this adds the task to processed
+                    self.finished_generations.append(result["generation"])
+                    if result["answer"] is not None:
+                        self.answer = json.loads(result["answer"])
+                        print("The answer and confidence after json parsing", self.answer)
+                        yield self.answer
+                    else:
+                        self.answer = await self.answerobj.extract_answer(self.question, result["generation"])
+                        print("The answer and confidence after interaction with 2nd model",self.answer)
+                        yield self.answer
+                except GeneratorExit:
+                    return
+                except Exception as e:
+                    traceback.print_exc()
+                    error_type = type(e).__name__
+                    print(f"[ERROR] {error_type}")
+                    traceback.print_exc()
+                    self.answer = {
+                        "Answer": -1,
+                        "Confidence": -0.1,
+                    }
+                    yield self.answer
+
+        finally:
+ #fallback in the Pipeline timeout handler. Timout
+            for t in self.tasks:
+                if t.done() and t not in processed:
+                    try:
+                        if not t.cancelled() and t.exception() is None:
+                            self.res = t.result()
+
+                        elif t.exception() is not None:
+                         # optional: handle failed tasks
+                         pass
+                    except Exception:
+                        pass
+                elif not t.done():
+                    t.cancel()
+            asyncio.create_task(abort_request(self.request_ids))
+
+            # Fire server-side abort independently — survives parent cancellation
+
+    def flatten_prompt_list(self):
+        global max_batch_size
+        self.flattened_prompt_list = [
+            self.system_prompt
+           # for system_prompt in self.prompts_list
+            for _ in range(max_batch_size)
+        ]
+
+    def generate_question_copies(self, question):
+        self.question = question
+        self.list_of_questions = [
+            prompt_template.fill(
+                input_dict={
+                    "problem": question,
+                    "system_prompt": system_prompt,
+                },
+                chat_template_kwargs = chat_template_kwargs,
+                format_as_string=True
+            )
+            for system_prompt in self.flattened_prompt_list
+        ]
+
+
+    async def predict_for_question(self, question):
+        self.flatten_prompt_list()
+        self.generate_question_copies(question)
+
+        gen = self.send_request_to_server()
+
+        try:
+            async for answer in gen:
+                yield answer
+
+        except Exception as e:
+            print("Error in predict_for_question:", e)
+            raise
+
+        finally:
+            try:
+                await gen.aclose()
+            except Exception:
+                pass
+
+import math
+
+class BufferBorrower:
+    """
+    Dynamic buffer-time borrowing strategy for inference.
+
+    Borrows from buffer time based on task difficulty and step-back
+    token usage, using a sigmoid curve for smooth allocation.
+
+    Parameters
+    ----------
+    max_difficulty : int or float
+        The upper bound of the difficulty scale (e.g., 5 or 1.0).
+    alpha : float
+        Weight for the difficulty signal (default 0.6).
+    beta : float
+        Weight for the step-back token signal (default 0.4).
+    b_max : float
+        Maximum fraction of buffer that can be borrowed (default 0.7).
+    k : float
+        Steepness of the sigmoid transition (default 6).
+    threshold : float
+        Midpoint of the sigmoid curve (default 0.4).
+    """
+
+    def __init__(
+        self,
+        b_max: float = 0.85,
+        k: float = 6.0,
+        threshold: float = 0.4,
+        total_questions: int = 50,
+        total_available_time: int = 15720,
+    ):
+
+        self.b_max = b_max
+        self.k = k
+        self.threshold = threshold
+        self.total_questions = total_questions
+        self.total_available_time = total_available_time
+
+    def compute_time_pressure(
+        self,
+        remaining_time: float,
+        questions_completed: int,
+        global_buffer: float = 0.0,
+    ) -> float:
+        remaining_q = max(1, self.total_questions - questions_completed)
+        if remaining_time <= 0:
+            return 1.5
+        ideal_pace = self.total_available_time / self.total_questions
+        available_pace = remaining_time / remaining_q
+        pressure = ideal_pace / available_pace
+        return max(0.3, min(1.5, pressure))
+
+    def allocate_time(
+        self,
+        remaining_time: float,
+        questions_completed: int,
+        global_buffer: float = 0.0,
+        allowed_time : float = 320,
+      ) -> dict:
+        """
+        Allocate effective inference and remaining buffer time.
+
+        Parameters
+        ----------
+        allowed_time : float
+            Base inference time budget.
+        global_buffer : float
+            global buffer time budget.
+        difficulty : float
+            Task difficulty score.
+        stepback_tokens : int
+            Tokens used in step-back phase.
+        stepback_budget : int
+            Total step-back token budget.
+
+        Returns
+        -------
+        dict
+            Keys: effective_inference, remaining_buffer, borrowed,
+                  borrow_fraction.
+        """
+        pressure = self.compute_time_pressure(
+                              remaining_time,
+                              questions_completed,
+                              global_buffer
+                            )
+        borrow_fraction = 1/pressure
+        max_borrowable = 130
+        print("borrow fraction", borrow_fraction)
+        borrowed = min(pressure * global_buffer, max_borrowable)
+
+
+        return {
+            "effective_inference": allowed_time + borrowed,
+            "global_buffer": global_buffer - borrowed,
+            "borrowed": borrowed,
+            "borrow_fraction": borrow_fraction,
+        }
+
+class TimeBudget:
+    def __init__(self, total_seconds):
+        self.start = time.perf_counter()
+        self.deadline = self.start + total_seconds
+
+    @property
+    def remaining(self):
+        return max(0, self.deadline - time.perf_counter())
+
+    @property
+    def elapsed(self):
+        return time.perf_counter() - self.start
+
+    @property
+    def expired(self):
+        return self.remaining <= 0
+
+class Pipeline:
+    def __init__(self):
+        self.budget_seconds = 0
+        self.k = 1
+        self.budget_seconds = 0
+    async def get_prediction(self, problem_text):
+        global global_buffer, i, borrower, max_batch_size,last_30, sampling_param
+        budgetobj = None
+        timeout = 60
+        # Timeout at this level - see if needs to be implemented
+        thresh_hold = 3
+        num_done = 0
+        max_generation_count = self.k*max_batch_size
+        answer_list = []
+        finalanswerobj = Result()
+        print("Pipeline step 1")
+        deadline = 0
+        allowed_time = 320
+        self.budget_seconds = allowed_time
+        if global_buffer> 0:
+            result = borrower.allocate_time(
+                    remaining_time = get_global_remaining(),
+                    questions_completed = i,
+                    allowed_time = allowed_time,
+                    global_buffer = global_buffer
+             )
+
+            self.budget_seconds = result["effective_inference"]
+            global_buffer = result["global_buffer"]
+            print(f'borrowed={result["borrowed"]:.0f}')
+        print(f"Budget: base={allowed_time:.0f}s "
+                          f"= {self.budget_seconds:.0f}s (global remaining: {get_global_remaining():.0f}s)")
+        budgetobj = TimeBudget(self.budget_seconds)
+
+        clientobj = ClientClass(default_prompt)
+        deadline = max(deadline, budgetobj.remaining)
+        operation_start_time = time.perf_counter()
+        print("Deadline is", deadline)
+        gen = clientobj.predict_for_question(problem_text)
+        try:
+            async with asyncio.timeout(deadline):
+                async for answer in gen:
+                    answer_list.append(answer)
+                    print("Answer list on timeout is:-")
+                    print(answer_list)
+                    num_done = len(answer_list)
+                    if num_done >= thresh_hold and num_done < max_generation_count:
+                        prediction, early_stop_flag = finalanswerobj.get_best_answer(answer_list, num_done, True)
+                        if early_stop_flag:
+                            return prediction
+
+                    elif num_done == max_generation_count:
+                        prediction, _ = finalanswerobj.get_best_answer(answer_list, num_done, False)
+                        return prediction
+                    else:
+                        continue
+        except (TimeoutError, asyncio.TimeoutError):
+            traceback.print_exc()
+            prediction, _ = finalanswerobj.get_best_answer(answer_list, num_done, False)
+            return prediction
+
+        except Exception as e:
+            traceback.print_exc()
+            print(f"UNEXPECTED ERROR: {type(e).__name__} {e}")
+            if answer_list:
+                prediction, _ = finalanswerobj.get_best_answer(answer_list, num_done, False)
+                return prediction
+            return None
+
+        finally:
+            await gen.aclose()
+            print("Operation duration", time.perf_counter()-operation_start_time)
+            if budgetobj.elapsed > self.budget_seconds:
+                global_buffer -= (budgetobj.elapsed - self.budget_seconds)
+            else:
+                global_buffer += (self.budget_seconds - budgetobj.elapsed)
+
+def predict(id_: pl.Series, problem: pl.Series) -> pl.DataFrame | pd.DataFrame:
+    """Make a prediction."""
+    global server_started, i
+    start_pred_time = time.perf_counter()
+    pipelineobj = Pipeline()
+    if server_started is False:
+        server_started = wait_for_server()
+
+    id_ = id_.item(0)
+    problem_text: str = problem.item(0)
+
+    # BUG FIX: compare duration to duration (was comparing duration to absolute timestamp)
+    if get_global_remaining() < 30:
+        return pl.DataFrame({"id": id_, "answer": 29443})
+    loop = asyncio.get_event_loop()
+    prediction = loop.run_until_complete(pipelineobj.get_prediction(problem_text))
+
+    # If prediction is still None after everything, use fallback
+    if prediction is None:
+        prediction = 29443
+
+    i = i + 1
+
+    print("Returned dataframe is ", pl.DataFrame({"id": id_, "answer": prediction}))
+    return pl.DataFrame({"id": id_, "answer": prediction})
+
+borrower = ""
+
+def run_local_inference(file_path: str, output_path: str = "submission.csv"):
+    global borrower
+    import pandas as pd
+    import polars as pl
+    borrower = BufferBorrower(total_questions = 50, total_available_time = get_global_remaining())
+    # Load file
+    if file_path.endswith(".xlsx"):
+        df = pd.read_excel(file_path)
+    else:
+        df = pd.read_csv(file_path)
+
+    # Basic validation
+    assert "problem" in df.columns, "Column 'problem' is required"
+    df = df.dropna(subset=["problem"])
+
+    # Optional: also remove rows where problem is just whitespace
+    df = df[df["problem"].str.strip() != ""]
+
+    if "id" not in df.columns:
+        df["id"] = range(len(df))
+
+    results = []
+
+    for idx, row in df.iterrows():
+        id_val = row["id"]
+        problem_text = row["problem"]
+
+        # Convert to polars Series (since your predict expects that)
+        id_series = pl.Series([id_val])
+        problem_series = pl.Series([problem_text])
+
+        try:
+            pred_df = predict(id_series, problem_series)
+
+            if isinstance(pred_df, pl.DataFrame):
+                pred = pred_df.to_pandas()
+            else:
+                pred = pred_df
+
+            results.append(pred.iloc[0])
+
+        except Exception as e:
+            print(f"Error at row {idx}: {e}")
+            results.append({"id": id_val, "answer": 29443})
+
+    final_df = pd.DataFrame(results)
+    final_df.to_csv(output_path, index=False)
+
+    print(f"✅ Saved predictions to {output_path}")
+    return final_df
+
+run_local_inference("/content/AIMO_ReferenceProblems.xlsx")