Function get_llm_response

scripts/eval_aime_benchmark.py:288–364

Get response from the LLM for a given problem. If multiple choices are returned, formats them as attempt dictionaries. Args: problem (str): The problem text model (str): The model identifier analyze_logits (bool): Whether to request logprobs Ret

(problem: str, model: str, analyze_logits: bool = False, extra_body: dict = None)

    }

def get_llm_response(problem: str, model: str, analyze_logits: bool = False, extra_body: dict = None) -> Union[str, List[Dict]]:
    """
    Get response from the LLM for a given problem.
    If multiple choices are returned, formats them as attempt dictionaries.
    
    Args:
        problem (str): The problem text
        model (str): The model identifier
        analyze_logits (bool): Whether to request logprobs
        
    Returns:
        Union[str, List[Dict]]: Either a string response or list of attempt dictionaries
    """
    try:
        # Add logprobs parameters if requested
        kwargs = {}
        if analyze_logits:
            kwargs["logprobs"] = True
            kwargs["top_logprobs"] = 3
        
        # Add extra_body if provided
        if extra_body:
            kwargs["extra_body"] = extra_body
        
        response = client.with_options(timeout=6000.0).chat.completions.create(
            model=model,
            messages=[
                {"role": "user", "content": SYSTEM_PROMPT + problem}
            ],
            max_tokens=64000,
            **kwargs
        )
        
        # Save raw response if logprobs are requested
        if analyze_logits:
            raw_filename = f"results/raw_responses_{model.replace('/', '_')}.json"
            problem_id = hash(problem) % 10000  # Simple hash to identify the problem
            save_raw_response(raw_filename, problem_id, response.model_dump())
        
        # If there's more than one choice, format as attempts
        if len(response.choices) > 1:
            attempts = []
            for i, choice in enumerate(response.choices):
                response_text = choice.message.content.strip()
                predicted_answer = extract_answer(response_text)
                attempt_data = {
                    "attempt_number": i + 1,
                    "response": response_text,
                    "predicted_answer": predicted_answer
                }
                
                # Add logprobs if available
                if analyze_logits and hasattr(choice.message, 'logprobs') and choice.message.logprobs:
                    attempt_data["logprobs"] = choice.message.logprobs
                
                attempts.append(attempt_data)
            return attempts