""" 评估器模块 — 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]