utils.py

import os
import re
import string
import random
from collections import Counter
from typing import Optional

import numpy as np
import torch


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)


def auto_device(device: Optional[str] = None) -> torch.device:
    if device is not None:
        return torch.device(device)
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


# ── HotpotQA answer normalisation (follows the official evaluation script) ──

def normalize_answer(s: Optional[str]) -> str:
    """Lower text, remove punctuation, articles, and extra whitespace."""
    if s is None:
        return ""

    def remove_articles(text):
        return re.sub(r"\b(a|an|the)\b", " ", text)

    def white_space_fix(text):
        return " ".join(text.split())

    def remove_punc(text):
        exclude = set(string.punctuation)
        return "".join(ch for ch in text if ch not in exclude)

    return white_space_fix(remove_articles(remove_punc(s.lower())))


def exact_match(prediction: str, gold: str) -> bool:
    return normalize_answer(prediction) == normalize_answer(gold)


def f1_score(prediction: str, gold: str) -> float:
    pred_tokens = normalize_answer(prediction).split()
    gold_tokens = normalize_answer(gold).split()
    common = Counter(pred_tokens) & Counter(gold_tokens)
    num_same = sum(common.values())
    if num_same == 0:
        return 0.0
    precision = num_same / len(pred_tokens)
    recall = num_same / len(gold_tokens)
    return 2 * precision * recall / (precision + recall)


def strip_think_tags(text: str) -> str:
    """
    Remove <think>...</think> blocks from Qwen3 output.
    Returns only the text after the closing </think> tag.
    Internal reasoning is not part of the communicated message.
    """
    # Remove complete <think>...</think> blocks
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    return text.strip()


def extract_answer(text: str) -> str:
    """
    Extract the final answer from model output.
    Looks for 'Final Answer: ...' pattern first, then falls back to
    the last non-empty line.
    """
    # Pattern 1: explicit "Final Answer:" marker
    match = re.search(r"[Ff]inal [Aa]nswer\s*[:\-]\s*(.+)", text)
    if match:
        return match.group(1).strip().rstrip(".")

    # Pattern 2: \boxed{...} (in case model uses math formatting)
    boxes = re.findall(r"\\boxed\{([^}]*)\}", text)
    if boxes:
        return boxes[-1].strip()

    # Pattern 3: last non-empty line
    lines = [ln.strip() for ln in text.strip().splitlines() if ln.strip()]
    if lines:
        return lines[-1]

    return ""