src/easytranslate/evaluation/evaluator.py

"""
评估器模块 — Person D 负责实现

功能要求:
将解码和评估指标整合为统一的评估接口。

使用方法:
    evaluator = Evaluator(model, tokenizer, config)
    results = evaluator.evaluate(test_loader)
"""

from __future__ import annotations

import logging
from typing import Optional

import torch
import torch.nn as nn
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader
from tqdm import tqdm

from easytranslate.evaluation.metrics import compute_all_metrics
from easytranslate.evaluation.decoding import greedy_decode, beam_search_decode, sample_decode

logger = logging.getLogger(__name__)


def _get_config_value(config, key_path, default=None):
    if config is None:
        return default
    if isinstance(config, dict):
        value = config
        for key in key_path:
            value = value.get(key, default)
            if value is default:
                break
        return value
    value = config
    for key in key_path:
        value = getattr(value, key, default)
        if value is default:
            break
    return value


class Evaluator:
    """
    翻译模型评估器。
    """

    def __init__(self, model: nn.Module, tokenizer, config: dict):
        self.model = model
        self.tokenizer = tokenizer
        self.config = config or {}

        self.evaluation_config = _get_config_value(self.config, ["evaluation"], {})
        self.decoding_config = _get_config_value(self.evaluation_config, ["decoding"], {})
        self.metrics = _get_config_value(self.evaluation_config, ["metrics"], ["bleu", "comet", "chrf", "ter"])
        self.strategy = self.decoding_config.get("strategy", "beam_search") if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "strategy", "beam_search")

        self.bos_id = self.tokenizer.bos_token_id
        self.eos_id = self.tokenizer.eos_token_id
        self.pad_id = self.tokenizer.pad_token_id

        self.max_decode_len = self.decoding_config.get("max_decode_len", 256) if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "max_decode_len", 256)
        self.beam_size = self.decoding_config.get("beam_size", 5) if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "beam_size", 5)
        self.length_penalty = self.decoding_config.get("length_penalty", 1.0) if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "length_penalty", 1.0)
        self.no_repeat_ngram_size = self.decoding_config.get("no_repeat_ngram_size", 0) if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "no_repeat_ngram_size", 0)
        self.temperature = self.decoding_config.get("sampling", {}).get("temperature", 1.0) if isinstance(self.decoding_config, dict) else getattr(getattr(self.decoding_config, "sampling", {}), "temperature", 1.0)
        self.top_k = self.decoding_config.get("sampling", {}).get("top_k", 0) if isinstance(self.decoding_config, dict) else getattr(getattr(self.decoding_config, "sampling", {}), "top_k", 0)
        self.top_p = self.decoding_config.get("sampling", {}).get("top_p", 1.0) if isinstance(self.decoding_config, dict) else getattr(getattr(self.decoding_config, "sampling", {}), "top_p", 1.0)

        self.use_generate = hasattr(self.model, "generate") and not (hasattr(self.model, "encode") and hasattr(self.model, "decode_step"))

    def _decode(self, src_ids: torch.Tensor, src_padding_mask: torch.BoolTensor) -> torch.Tensor:
        if self.use_generate:
            generate_kwargs = {
                "max_length": self.max_decode_len,
                "early_stopping": True,
            }
            if self.strategy == "beam_search":
                generate_kwargs.update(
                    {
                        "num_beams": self.beam_size,
                        "length_penalty": self.length_penalty,
                        "no_repeat_ngram_size": self.no_repeat_ngram_size,
                    }
                )
            elif self.strategy == "sampling":
                generate_kwargs.update(
                    {
                        "do_sample": True,
                        "temperature": self.temperature,
                        "top_k": self.top_k,
                        "top_p": self.top_p,
                        "num_beams": 1,
                    }
                )
            else:
                generate_kwargs.update({"num_beams": 1})

            attention_mask = (~src_padding_mask).long()
            return self.model.generate(input_ids=src_ids, attention_mask=attention_mask, **generate_kwargs)

        if self.strategy == "beam_search":
            return beam_search_decode(
                self.model,
                src_ids,
                src_padding_mask,
                self.bos_id,
                self.eos_id,
                beam_size=self.beam_size,
                max_len=self.max_decode_len,
                length_penalty=self.length_penalty,
                no_repeat_ngram_size=self.no_repeat_ngram_size,
            )
        if self.strategy == "sampling":
            return sample_decode(
                self.model,
                src_ids,
                src_padding_mask,
                self.bos_id,
                self.eos_id,
                max_len=self.max_decode_len,
                temperature=self.temperature,
                top_k=self.top_k,
                top_p=self.top_p,
            )
        return greedy_decode(
            self.model,
            src_ids,
            src_padding_mask,
            self.bos_id,
            self.eos_id,
            max_len=self.max_decode_len,
        )

    def evaluate(
        self,
        dataloader: DataLoader,
        src_texts: Optional[list[str]] = None,
        ref_texts: Optional[list[str]] = None,
    ) -> dict:
        self.model.eval()
        device = next(self.model.parameters()).device if any(p.requires_grad or p.is_floating_point() for p in self.model.parameters()) else torch.device("cpu")

        if src_texts is None and hasattr(dataloader, "dataset") and hasattr(dataloader.dataset, "src_texts"):
            src_texts = list(dataloader.dataset.src_texts)
        if ref_texts is None and hasattr(dataloader, "dataset") and hasattr(dataloader.dataset, "tgt_texts"):
            ref_texts = list(dataloader.dataset.tgt_texts)

        if src_texts is None or ref_texts is None:
            raise ValueError("Source texts and reference texts must be provided for evaluation.")

        hypotheses: list[str] = []
        sources: list[str] = []
        references: list[str] = []

        for batch_idx, batch in enumerate(tqdm(dataloader, desc="Evaluating", unit="batch")):
            src_ids = batch["src_ids"].to(device)
            src_padding_mask = batch.get("src_padding_mask")
            if src_padding_mask is None:
                src_padding_mask = src_ids.eq(self.pad_id)
            else:
                src_padding_mask = src_padding_mask.to(device)

            output_ids = self._decode(src_ids, src_padding_mask)
            if isinstance(output_ids, torch.Tensor):
                output_ids = output_ids.cpu()

            for sample_idx in range(output_ids.size(0)):
                decoded = self.tokenizer.decode(output_ids[sample_idx].tolist(), skip_special_tokens=True)
                hypotheses.append(decoded)

        sources = src_texts
        references = ref_texts
        results = compute_all_metrics(sources, hypotheses, references, metrics=self.metrics)
        return results

    def translate(self, texts: list[str]) -> list[str]:
        self.model.eval()
        device = next(self.model.parameters()).device if any(p.requires_grad or p.is_floating_point() for p in self.model.parameters()) else torch.device("cpu")

        input_ids = []
        for text in texts:
            src_tokens = self.tokenizer.encode(
                text,
                add_special_tokens=True,
                max_length=_get_config_value(self.config, ["data", "preprocessing", "max_src_len"], 256),
            )
            input_ids.append(torch.tensor(src_tokens, dtype=torch.long, device=device))

        src_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.pad_id)
        src_padding_mask = src_ids.eq(self.pad_id)

        output_ids = self._decode(src_ids, src_padding_mask)
        if isinstance(output_ids, torch.Tensor):
            output_ids = output_ids.cpu()

        translations: list[str] = []
        for sample_idx in range(output_ids.size(0)):
            translations.append(self.tokenizer.decode(output_ids[sample_idx].tolist(), skip_special_tokens=True))

        return translations

    def translate_single(self, text: str) -> str:
        return self.translate([text])[0]