[Person D] Implement evaluation module and inference scripts
Browse files## Summary
Implement the evaluation and inference module (Person D's assignment):
### Evaluation metrics (`src/easytranslate/evaluation/metrics.py`)
- SacreBLEU with Chinese tokenization
- COMET neural evaluation metric
- chrF++ and TER metrics
### Decoding strategies (`src/easytranslate/evaluation/decoding.py`)
- Greedy decoding
- Beam search with length penalty
- Sampling with temperature, top-k, and top-p
### Evaluator (`src/easytranslate/evaluation/evaluator.py`)
- Unified evaluation interface
- Batch evaluation support
### Scripts
- `scripts/evaluate.py`: Evaluation entry point
- `scripts/translate.py`: CLI interactive translation + file translation
- scripts/evaluate.py +285 -50
- scripts/translate.py +316 -82
- src/easytranslate/evaluation/__init__.py +16 -16
- src/easytranslate/evaluation/decoding.py +219 -129
- src/easytranslate/evaluation/evaluator.py +208 -76
- src/easytranslate/evaluation/metrics.py +169 -109
scripts/evaluate.py
CHANGED
|
@@ -1,50 +1,285 @@
|
|
| 1 |
-
"""
|
| 2 |
-
评估入口脚本
|
| 3 |
-
|
| 4 |
-
使用方式:
|
| 5 |
-
# 在测试集上评估
|
| 6 |
-
python scripts/evaluate.py --config configs/default_config.yaml --checkpoint checkpoints/best_model.pt
|
| 7 |
-
|
| 8 |
-
# 指定解码策略
|
| 9 |
-
python scripts/evaluate.py --checkpoint checkpoints/best_model.pt evaluation.decoding.strategy=beam_search evaluation.decoding.beam_size=10
|
| 10 |
-
"""
|
| 11 |
-
|
| 12 |
-
import argparse
|
| 13 |
-
import
|
| 14 |
-
|
| 15 |
-
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
|
| 31 |
-
|
| 32 |
-
|
| 33 |
-
|
| 34 |
-
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
|
| 46 |
-
|
| 47 |
-
|
| 48 |
-
|
| 49 |
-
|
| 50 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
评估入口脚本
|
| 3 |
+
|
| 4 |
+
使用方式:
|
| 5 |
+
# 在测试集上评估
|
| 6 |
+
python scripts/evaluate.py --config configs/default_config.yaml --checkpoint checkpoints/best_model.pt
|
| 7 |
+
|
| 8 |
+
# 指定解码策略
|
| 9 |
+
python scripts/evaluate.py --checkpoint checkpoints/best_model.pt evaluation.decoding.strategy=beam_search evaluation.decoding.beam_size=10
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
import argparse
|
| 13 |
+
import json
|
| 14 |
+
import sys
|
| 15 |
+
from pathlib import Path
|
| 16 |
+
|
| 17 |
+
try:
|
| 18 |
+
from omegaconf import OmegaConf
|
| 19 |
+
except ImportError: # pragma: no cover
|
| 20 |
+
OmegaConf = None
|
| 21 |
+
import yaml
|
| 22 |
+
|
| 23 |
+
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
|
| 24 |
+
|
| 25 |
+
import torch
|
| 26 |
+
from torch.utils.data import DataLoader
|
| 27 |
+
|
| 28 |
+
from easytranslate.data.collator import TranslationCollator
|
| 29 |
+
from easytranslate.data.dataset import (
|
| 30 |
+
TranslationDataset,
|
| 31 |
+
load_custom_dataset,
|
| 32 |
+
load_opus_dataset,
|
| 33 |
+
load_wmt_dataset,
|
| 34 |
+
)
|
| 35 |
+
from easytranslate.data.tokenizer import TokenizerWrapper, build_tokenizer
|
| 36 |
+
from easytranslate.evaluation.evaluator import Evaluator
|
| 37 |
+
from easytranslate.model import TransformerTranslationModel
|
| 38 |
+
from easytranslate.model.finetune import load_pretrained_model
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def parse_args():
|
| 42 |
+
parser = argparse.ArgumentParser(description="EasyTranslate Evaluation")
|
| 43 |
+
parser.add_argument("--config", type=str, default="configs/default_config.yaml")
|
| 44 |
+
parser.add_argument("--checkpoint", type=str, required=True, help="模型检查点路径")
|
| 45 |
+
parser.add_argument("--output", type=str, default="outputs/evaluation_results.json", help="结果保存路径")
|
| 46 |
+
args, unknown = parser.parse_known_args()
|
| 47 |
+
return args, unknown
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def _get_config(config, *keys, default=None):
|
| 51 |
+
value = config
|
| 52 |
+
for key in keys:
|
| 53 |
+
if isinstance(value, dict):
|
| 54 |
+
value = value.get(key, default)
|
| 55 |
+
else:
|
| 56 |
+
value = getattr(value, key, default)
|
| 57 |
+
if value is default:
|
| 58 |
+
break
|
| 59 |
+
return value
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def _load_config(path, cli_overrides=None):
|
| 63 |
+
if OmegaConf is not None:
|
| 64 |
+
config = OmegaConf.load(path)
|
| 65 |
+
if cli_overrides:
|
| 66 |
+
config = OmegaConf.merge(config, OmegaConf.from_cli(cli_overrides))
|
| 67 |
+
return config
|
| 68 |
+
|
| 69 |
+
with open(path, "r", encoding="utf-8") as fin:
|
| 70 |
+
config = yaml.safe_load(fin)
|
| 71 |
+
if cli_overrides:
|
| 72 |
+
print("Warning: OmegaConf is not installed; CLI overrides are ignored.")
|
| 73 |
+
return config
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def load_test_split(config):
|
| 77 |
+
dataset_name = _get_config(config, "data", "dataset_name")
|
| 78 |
+
|
| 79 |
+
if dataset_name == "wmt":
|
| 80 |
+
try:
|
| 81 |
+
dataset = load_wmt_dataset(
|
| 82 |
+
year=_get_config(config, "data", "wmt", "year"),
|
| 83 |
+
language_pair=_get_config(config, "data", "wmt", "language_pair"),
|
| 84 |
+
split="test",
|
| 85 |
+
)
|
| 86 |
+
except Exception:
|
| 87 |
+
dataset = load_wmt_dataset(
|
| 88 |
+
year=_get_config(config, "data", "wmt", "year"),
|
| 89 |
+
language_pair=_get_config(config, "data", "wmt", "language_pair"),
|
| 90 |
+
split="validation",
|
| 91 |
+
)
|
| 92 |
+
return list(dataset["src"]), list(dataset["tgt"])
|
| 93 |
+
|
| 94 |
+
if dataset_name == "opus":
|
| 95 |
+
try:
|
| 96 |
+
dataset = load_opus_dataset(
|
| 97 |
+
subset=_get_config(config, "data", "opus", "subset"),
|
| 98 |
+
split="test",
|
| 99 |
+
)
|
| 100 |
+
except Exception:
|
| 101 |
+
dataset = load_opus_dataset(
|
| 102 |
+
subset=_get_config(config, "data", "opus", "subset"),
|
| 103 |
+
split="validation",
|
| 104 |
+
)
|
| 105 |
+
return list(dataset["src"]), list(dataset["tgt"])
|
| 106 |
+
|
| 107 |
+
if dataset_name == "custom":
|
| 108 |
+
data = load_custom_dataset(
|
| 109 |
+
train_src=_get_config(config, "data", "custom", "train_src"),
|
| 110 |
+
train_tgt=_get_config(config, "data", "custom", "train_tgt"),
|
| 111 |
+
val_src=_get_config(config, "data", "custom", "val_src"),
|
| 112 |
+
val_tgt=_get_config(config, "data", "custom", "val_tgt"),
|
| 113 |
+
test_src=_get_config(config, "data", "custom", "test_src"),
|
| 114 |
+
test_tgt=_get_config(config, "data", "custom", "test_tgt"),
|
| 115 |
+
preprocessing_config=_get_config(config, "data", "preprocessing"),
|
| 116 |
+
)
|
| 117 |
+
if "test" not in data:
|
| 118 |
+
raise ValueError("Custom dataset missing test split")
|
| 119 |
+
return data["test"]["src"], data["test"]["tgt"]
|
| 120 |
+
|
| 121 |
+
raise ValueError(f"Unsupported dataset_name: {dataset_name}")
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
def _get_tokenizer_train_texts(config, allow_auto: bool = False) -> list[str] | None:
|
| 125 |
+
dataset_name = _get_config(config, "data", "dataset_name")
|
| 126 |
+
if dataset_name == "custom":
|
| 127 |
+
custom = _get_config(config, "data", "custom") or {}
|
| 128 |
+
data = load_custom_dataset(
|
| 129 |
+
train_src=custom.get("train_src"),
|
| 130 |
+
train_tgt=custom.get("train_tgt"),
|
| 131 |
+
val_src=custom.get("val_src"),
|
| 132 |
+
val_tgt=custom.get("val_tgt"),
|
| 133 |
+
test_src=custom.get("test_src"),
|
| 134 |
+
test_tgt=custom.get("test_tgt"),
|
| 135 |
+
preprocessing_config=_get_config(config, "data", "preprocessing"),
|
| 136 |
+
)
|
| 137 |
+
return list(data["train"]["src"]) + list(data["train"]["tgt"])
|
| 138 |
+
|
| 139 |
+
if not allow_auto:
|
| 140 |
+
return None
|
| 141 |
+
|
| 142 |
+
if dataset_name == "wmt":
|
| 143 |
+
try:
|
| 144 |
+
dataset = load_wmt_dataset(
|
| 145 |
+
year=_get_config(config, "data", "wmt", "year"),
|
| 146 |
+
language_pair=_get_config(config, "data", "wmt", "language_pair"),
|
| 147 |
+
split="train",
|
| 148 |
+
)
|
| 149 |
+
except Exception:
|
| 150 |
+
dataset = load_wmt_dataset(
|
| 151 |
+
year=_get_config(config, "data", "wmt", "year"),
|
| 152 |
+
language_pair=_get_config(config, "data", "wmt", "language_pair"),
|
| 153 |
+
split="validation",
|
| 154 |
+
)
|
| 155 |
+
return list(dataset["src"]) + list(dataset["tgt"])
|
| 156 |
+
|
| 157 |
+
if dataset_name == "opus":
|
| 158 |
+
try:
|
| 159 |
+
dataset = load_opus_dataset(
|
| 160 |
+
subset=_get_config(config, "data", "opus", "subset"),
|
| 161 |
+
split="train",
|
| 162 |
+
)
|
| 163 |
+
except Exception:
|
| 164 |
+
dataset = load_opus_dataset(
|
| 165 |
+
subset=_get_config(config, "data", "opus", "subset"),
|
| 166 |
+
split="validation",
|
| 167 |
+
)
|
| 168 |
+
return list(dataset["src"]) + list(dataset["tgt"])
|
| 169 |
+
|
| 170 |
+
return None
|
| 171 |
+
|
| 172 |
+
|
| 173 |
+
def build_model_and_tokenizer(config, device):
|
| 174 |
+
model_type = _get_config(config, "model", "type")
|
| 175 |
+
if model_type == "transformer_scratch":
|
| 176 |
+
tokenizer_config = _get_config(config, "tokenizer") or {}
|
| 177 |
+
tokenizer_path = tokenizer_config.get("path") or tokenizer_config.get("tokenizer_path")
|
| 178 |
+
tokenizer_type = tokenizer_config.get("type", "bpe")
|
| 179 |
+
auto_train = bool(tokenizer_config.get("auto_train", False))
|
| 180 |
+
if tokenizer_type in {"bpe", "sentencepiece"} and not tokenizer_path:
|
| 181 |
+
train_texts = _get_tokenizer_train_texts(config, allow_auto=auto_train)
|
| 182 |
+
if train_texts is None:
|
| 183 |
+
raise ValueError(
|
| 184 |
+
"BPE tokenizer requires tokenizer.path or a local custom dataset with train texts. "
|
| 185 |
+
"Automatic WMT/OPUS download is disabled by default. "
|
| 186 |
+
"Set tokenizer.auto_train=true to enable it, or provide tokenizer.path/pretrained tokenizer."
|
| 187 |
+
)
|
| 188 |
+
tokenizer = build_tokenizer(tokenizer_config, train_texts=train_texts)
|
| 189 |
+
else:
|
| 190 |
+
tokenizer = build_tokenizer(tokenizer_config)
|
| 191 |
+
model = TransformerTranslationModel(
|
| 192 |
+
src_vocab_size=tokenizer.vocab_size,
|
| 193 |
+
tgt_vocab_size=tokenizer.vocab_size,
|
| 194 |
+
d_model=_get_config(config, "model", "transformer", "d_model"),
|
| 195 |
+
nhead=_get_config(config, "model", "transformer", "nhead"),
|
| 196 |
+
num_encoder_layers=_get_config(config, "model", "transformer", "num_encoder_layers"),
|
| 197 |
+
num_decoder_layers=_get_config(config, "model", "transformer", "num_decoder_layers"),
|
| 198 |
+
dim_feedforward=_get_config(config, "model", "transformer", "dim_feedforward"),
|
| 199 |
+
dropout=_get_config(config, "model", "transformer", "dropout"),
|
| 200 |
+
activation=_get_config(config, "model", "transformer", "activation"),
|
| 201 |
+
max_seq_len=_get_config(config, "model", "transformer", "max_seq_len"),
|
| 202 |
+
use_flash_attention=_get_config(config, "model", "transformer", "use_flash_attention"),
|
| 203 |
+
use_rotary_embedding=_get_config(config, "model", "transformer", "use_rotary_embedding"),
|
| 204 |
+
pre_norm=_get_config(config, "model", "transformer", "pre_norm"),
|
| 205 |
+
pad_id=tokenizer.pad_token_id,
|
| 206 |
+
)
|
| 207 |
+
return model.to(device), tokenizer
|
| 208 |
+
|
| 209 |
+
model, hf_tokenizer = load_pretrained_model(
|
| 210 |
+
config.model.pretrained.model_name,
|
| 211 |
+
config.model.pretrained.src_lang,
|
| 212 |
+
config.model.pretrained.tgt_lang,
|
| 213 |
+
device=str(device),
|
| 214 |
+
)
|
| 215 |
+
tokenizer = TokenizerWrapper(
|
| 216 |
+
hf_tokenizer,
|
| 217 |
+
pad_token=getattr(hf_tokenizer, "pad_token", "<pad>"),
|
| 218 |
+
unk_token=getattr(hf_tokenizer, "unk_token", "<unk>"),
|
| 219 |
+
bos_token=getattr(hf_tokenizer, "bos_token", "<s>"),
|
| 220 |
+
eos_token=getattr(hf_tokenizer, "eos_token", "</s>"),
|
| 221 |
+
)
|
| 222 |
+
return model, tokenizer
|
| 223 |
+
|
| 224 |
+
|
| 225 |
+
def load_checkpoint(model, checkpoint_path, device):
|
| 226 |
+
checkpoint = torch.load(checkpoint_path, map_location=device)
|
| 227 |
+
if isinstance(checkpoint, dict):
|
| 228 |
+
if "model_state_dict" in checkpoint:
|
| 229 |
+
model.load_state_dict(checkpoint["model_state_dict"])
|
| 230 |
+
elif "state_dict" in checkpoint:
|
| 231 |
+
model.load_state_dict(checkpoint["state_dict"])
|
| 232 |
+
else:
|
| 233 |
+
try:
|
| 234 |
+
model.load_state_dict(checkpoint)
|
| 235 |
+
except Exception as exc:
|
| 236 |
+
raise ValueError("Checkpoint does not contain a valid model state dict") from exc
|
| 237 |
+
else:
|
| 238 |
+
raise ValueError("Unsupported checkpoint format")
|
| 239 |
+
return model
|
| 240 |
+
|
| 241 |
+
|
| 242 |
+
def main():
|
| 243 |
+
args, cli_overrides = parse_args()
|
| 244 |
+
|
| 245 |
+
print("=" * 60)
|
| 246 |
+
print(" EasyTranslate - Evaluation")
|
| 247 |
+
print("=" * 60)
|
| 248 |
+
|
| 249 |
+
config = _load_config(args.config, cli_overrides)
|
| 250 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 251 |
+
model, tokenizer = build_model_and_tokenizer(config, device)
|
| 252 |
+
model = load_checkpoint(model, args.checkpoint, device)
|
| 253 |
+
|
| 254 |
+
src_texts, tgt_texts = load_test_split(config)
|
| 255 |
+
|
| 256 |
+
dataset = TranslationDataset(
|
| 257 |
+
src_texts=src_texts,
|
| 258 |
+
tgt_texts=tgt_texts,
|
| 259 |
+
tokenizer=tokenizer,
|
| 260 |
+
max_src_len=_get_config(config, "data", "preprocessing", "max_src_len"),
|
| 261 |
+
max_tgt_len=_get_config(config, "data", "preprocessing", "max_tgt_len"),
|
| 262 |
+
)
|
| 263 |
+
collator = TranslationCollator(pad_token_id=tokenizer.pad_token_id)
|
| 264 |
+
dataloader = DataLoader(
|
| 265 |
+
dataset,
|
| 266 |
+
batch_size=config.data.dataloader.batch_size,
|
| 267 |
+
shuffle=False,
|
| 268 |
+
num_workers=int(config.data.dataloader.num_workers),
|
| 269 |
+
pin_memory=bool(config.data.dataloader.pin_memory),
|
| 270 |
+
collate_fn=collator,
|
| 271 |
+
)
|
| 272 |
+
|
| 273 |
+
evaluator = Evaluator(model, tokenizer, config)
|
| 274 |
+
results = evaluator.evaluate(dataloader, src_texts=src_texts, ref_texts=tgt_texts)
|
| 275 |
+
|
| 276 |
+
output_path = Path(args.output)
|
| 277 |
+
output_path.parent.mkdir(parents=True, exist_ok=True)
|
| 278 |
+
with output_path.open("w", encoding="utf-8") as fout:
|
| 279 |
+
json.dump(results, fout, ensure_ascii=False, indent=2)
|
| 280 |
+
|
| 281 |
+
print(json.dumps(results, ensure_ascii=False, indent=2))
|
| 282 |
+
|
| 283 |
+
|
| 284 |
+
if __name__ == "__main__":
|
| 285 |
+
main()
|
scripts/translate.py
CHANGED
|
@@ -1,82 +1,316 @@
|
|
| 1 |
-
"""
|
| 2 |
-
交互式翻译推理脚本
|
| 3 |
-
|
| 4 |
-
使用方式:
|
| 5 |
-
# 命令行交互翻译
|
| 6 |
-
python scripts/translate.py --checkpoint checkpoints/best_model.pt
|
| 7 |
-
|
| 8 |
-
# 翻译文件
|
| 9 |
-
python scripts/translate.py --checkpoint checkpoints/best_model.pt --input input.txt --output output.txt
|
| 10 |
-
|
| 11 |
-
# 启动
|
| 12 |
-
python scripts/translate.py --
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
|
| 31 |
-
|
| 32 |
-
|
| 33 |
-
|
| 34 |
-
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
def
|
| 46 |
-
""
|
| 47 |
-
|
| 48 |
-
|
| 49 |
-
|
| 50 |
-
|
| 51 |
-
|
| 52 |
-
|
| 53 |
-
|
| 54 |
-
|
| 55 |
-
|
| 56 |
-
|
| 57 |
-
def
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
|
| 61 |
-
|
| 62 |
-
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
|
| 67 |
-
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
|
| 71 |
-
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
|
| 75 |
-
|
| 76 |
-
|
| 77 |
-
|
| 78 |
-
|
| 79 |
-
|
| 80 |
-
|
| 81 |
-
|
| 82 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
交互式翻译推理脚本
|
| 3 |
+
|
| 4 |
+
使用方式:
|
| 5 |
+
# 命令行交互翻译
|
| 6 |
+
python scripts/translate.py --checkpoint checkpoints/best_model.pt
|
| 7 |
+
|
| 8 |
+
# 翻译文件
|
| 9 |
+
python scripts/translate.py --checkpoint checkpoints/best_model.pt --input input.txt --output output.txt
|
| 10 |
+
|
| 11 |
+
# 启动 Streamlit Web UI(无需模型即可预览界面)
|
| 12 |
+
python scripts/translate.py --web
|
| 13 |
+
# 或者带模型启动翻译
|
| 14 |
+
python scripts/translate.py --checkpoint checkpoints/best_model.pt --web
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
import argparse
|
| 18 |
+
import subprocess
|
| 19 |
+
import sys
|
| 20 |
+
from pathlib import Path
|
| 21 |
+
|
| 22 |
+
try:
|
| 23 |
+
from omegaconf import OmegaConf
|
| 24 |
+
except ImportError: # pragma: no cover
|
| 25 |
+
OmegaConf = None
|
| 26 |
+
import yaml
|
| 27 |
+
|
| 28 |
+
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
|
| 29 |
+
|
| 30 |
+
import torch
|
| 31 |
+
|
| 32 |
+
from easytranslate.data.collator import TranslationCollator
|
| 33 |
+
from easytranslate.data.dataset import (
|
| 34 |
+
TranslationDataset,
|
| 35 |
+
load_custom_dataset,
|
| 36 |
+
load_opus_dataset,
|
| 37 |
+
load_wmt_dataset,
|
| 38 |
+
)
|
| 39 |
+
from easytranslate.data.tokenizer import TokenizerWrapper, build_tokenizer
|
| 40 |
+
from easytranslate.evaluation.evaluator import Evaluator
|
| 41 |
+
from easytranslate.model import TransformerTranslationModel
|
| 42 |
+
from easytranslate.model.finetune import load_pretrained_model
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def parse_args():
|
| 46 |
+
parser = argparse.ArgumentParser(description="EasyTranslate Inference")
|
| 47 |
+
parser.add_argument("--config", type=str, default="configs/default_config.yaml")
|
| 48 |
+
parser.add_argument("--checkpoint", type=str, default=None)
|
| 49 |
+
parser.add_argument("--input", type=str, default=None, help="输入文件路径")
|
| 50 |
+
parser.add_argument("--output", type=str, default=None, help="输出文件路径")
|
| 51 |
+
parser.add_argument("--web", action="store_true", help="启动 Streamlit Web UI")
|
| 52 |
+
parser.add_argument("--streamlit-app", action="store_true", help=argparse.SUPPRESS)
|
| 53 |
+
args, unknown = parser.parse_known_args()
|
| 54 |
+
return args, unknown
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
def _get_config(config, *keys, default=None):
|
| 58 |
+
value = config
|
| 59 |
+
for key in keys:
|
| 60 |
+
if isinstance(value, dict):
|
| 61 |
+
value = value.get(key, default)
|
| 62 |
+
else:
|
| 63 |
+
value = getattr(value, key, default)
|
| 64 |
+
if value is default:
|
| 65 |
+
break
|
| 66 |
+
return value
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def _load_config(path, cli_overrides=None):
|
| 70 |
+
if OmegaConf is not None:
|
| 71 |
+
config = OmegaConf.load(path)
|
| 72 |
+
if cli_overrides:
|
| 73 |
+
config = OmegaConf.merge(config, OmegaConf.from_cli(cli_overrides))
|
| 74 |
+
return config
|
| 75 |
+
|
| 76 |
+
with open(path, "r", encoding="utf-8") as fin:
|
| 77 |
+
config = yaml.safe_load(fin)
|
| 78 |
+
if cli_overrides:
|
| 79 |
+
print("Warning: OmegaConf is not installed; CLI overrides are ignored.")
|
| 80 |
+
return config
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
def interactive_translate(evaluator):
|
| 84 |
+
print("输入英⽂句⼦,按回车翻译;输入 'quit' 退出。")
|
| 85 |
+
while True:
|
| 86 |
+
try:
|
| 87 |
+
text = input("> ").strip()
|
| 88 |
+
except EOFError:
|
| 89 |
+
break
|
| 90 |
+
if not text:
|
| 91 |
+
continue
|
| 92 |
+
if text.lower() in {"quit", "exit"}:
|
| 93 |
+
break
|
| 94 |
+
translation = evaluator.translate_single(text)
|
| 95 |
+
print(translation)
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
def translate_file(evaluator, input_path: str, output_path: str):
|
| 99 |
+
source_lines = []
|
| 100 |
+
with Path(input_path).open("r", encoding="utf-8") as fin:
|
| 101 |
+
for line in fin:
|
| 102 |
+
line = line.strip()
|
| 103 |
+
if line:
|
| 104 |
+
source_lines.append(line)
|
| 105 |
+
|
| 106 |
+
translations = evaluator.translate(source_lines)
|
| 107 |
+
|
| 108 |
+
output_file = Path(output_path)
|
| 109 |
+
output_file.parent.mkdir(parents=True, exist_ok=True)
|
| 110 |
+
with output_file.open("w", encoding="utf-8") as fout:
|
| 111 |
+
for line in translations:
|
| 112 |
+
fout.write(f"{line}\n")
|
| 113 |
+
|
| 114 |
+
print(f"Translation complete: {len(translations)} lines written to {output_path}")
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
def launch_streamlit_app(args):
|
| 118 |
+
import streamlit as st
|
| 119 |
+
|
| 120 |
+
@st.cache_resource
|
| 121 |
+
def load_evaluator():
|
| 122 |
+
config = _load_config(args.config)
|
| 123 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 124 |
+
model, tokenizer = build_model_and_tokenizer(config, device)
|
| 125 |
+
model = load_checkpoint(model, args.checkpoint, device)
|
| 126 |
+
return Evaluator(model, tokenizer, config)
|
| 127 |
+
|
| 128 |
+
st.set_page_config(page_title="EasyTranslate", layout="wide")
|
| 129 |
+
st.title("EasyTranslate")
|
| 130 |
+
st.write("English to Chinese translation powered by EasyTranslate.")
|
| 131 |
+
|
| 132 |
+
if args.checkpoint is None:
|
| 133 |
+
st.warning(
|
| 134 |
+
"当前未提供模型 checkpoint,页面仅用于预览界面效果。"
|
| 135 |
+
" 如需翻译,请传入 --checkpoint 或先训练生成模型。"
|
| 136 |
+
)
|
| 137 |
+
|
| 138 |
+
input_text = st.text_area("English Input", value="", height=200)
|
| 139 |
+
if st.button("Translate"):
|
| 140 |
+
if not input_text.strip():
|
| 141 |
+
st.warning("请输入要翻译的英文文本。")
|
| 142 |
+
elif args.checkpoint is None:
|
| 143 |
+
st.error("未提供 checkpoint,无法执行翻译。请使用 --checkpoint 参数启动。")
|
| 144 |
+
else:
|
| 145 |
+
with st.spinner("Translating..."):
|
| 146 |
+
try:
|
| 147 |
+
evaluator = load_evaluator()
|
| 148 |
+
translation = evaluator.translate_single(input_text.strip())
|
| 149 |
+
st.text_area("Chinese Translation", value=translation, height=200)
|
| 150 |
+
except Exception as exc:
|
| 151 |
+
st.error(f"模型加载或翻译失败:{exc}")
|
| 152 |
+
|
| 153 |
+
st.markdown("---")
|
| 154 |
+
st.caption("此页面用于展示前端界面;在未提供模型时,翻译功能会被禁用。")
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
def launch_streamlit_process(args):
|
| 158 |
+
script_path = Path(__file__).resolve()
|
| 159 |
+
cmd = [sys.executable, "-m", "streamlit", "run", str(script_path), "--", "--streamlit-app", "--config", args.config]
|
| 160 |
+
if args.checkpoint:
|
| 161 |
+
cmd.extend(["--checkpoint", args.checkpoint])
|
| 162 |
+
if args.input:
|
| 163 |
+
cmd.extend(["--input", args.input])
|
| 164 |
+
if args.output:
|
| 165 |
+
cmd.extend(["--output", args.output])
|
| 166 |
+
subprocess.run(cmd, check=True)
|
| 167 |
+
|
| 168 |
+
|
| 169 |
+
def _get_tokenizer_train_texts(config, allow_auto: bool = False) -> list[str] | None:
|
| 170 |
+
dataset_name = _get_config(config, "data", "dataset_name")
|
| 171 |
+
if dataset_name == "custom":
|
| 172 |
+
custom = _get_config(config, "data", "custom") or {}
|
| 173 |
+
data = load_custom_dataset(
|
| 174 |
+
train_src=custom.get("train_src"),
|
| 175 |
+
train_tgt=custom.get("train_tgt"),
|
| 176 |
+
val_src=custom.get("val_src"),
|
| 177 |
+
val_tgt=custom.get("val_tgt"),
|
| 178 |
+
test_src=custom.get("test_src"),
|
| 179 |
+
test_tgt=custom.get("test_tgt"),
|
| 180 |
+
preprocessing_config=_get_config(config, "data", "preprocessing"),
|
| 181 |
+
)
|
| 182 |
+
return list(data["train"]["src"]) + list(data["train"]["tgt"])
|
| 183 |
+
|
| 184 |
+
if not allow_auto:
|
| 185 |
+
return None
|
| 186 |
+
|
| 187 |
+
if dataset_name == "wmt":
|
| 188 |
+
try:
|
| 189 |
+
dataset = load_wmt_dataset(
|
| 190 |
+
year=_get_config(config, "data", "wmt", "year"),
|
| 191 |
+
language_pair=_get_config(config, "data", "wmt", "language_pair"),
|
| 192 |
+
split="train",
|
| 193 |
+
)
|
| 194 |
+
except Exception:
|
| 195 |
+
dataset = load_wmt_dataset(
|
| 196 |
+
year=_get_config(config, "data", "wmt", "year"),
|
| 197 |
+
language_pair=_get_config(config, "data", "wmt", "language_pair"),
|
| 198 |
+
split="validation",
|
| 199 |
+
)
|
| 200 |
+
return list(dataset["src"]) + list(dataset["tgt"])
|
| 201 |
+
|
| 202 |
+
if dataset_name == "opus":
|
| 203 |
+
try:
|
| 204 |
+
dataset = load_opus_dataset(
|
| 205 |
+
subset=_get_config(config, "data", "opus", "subset"),
|
| 206 |
+
split="train",
|
| 207 |
+
)
|
| 208 |
+
except Exception:
|
| 209 |
+
dataset = load_opus_dataset(
|
| 210 |
+
subset=_get_config(config, "data", "opus", "subset"),
|
| 211 |
+
split="validation",
|
| 212 |
+
)
|
| 213 |
+
return list(dataset["src"]) + list(dataset["tgt"])
|
| 214 |
+
|
| 215 |
+
return None
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
def build_model_and_tokenizer(config, device):
|
| 219 |
+
model_type = _get_config(config, "model", "type")
|
| 220 |
+
if model_type == "transformer_scratch":
|
| 221 |
+
tokenizer_config = _get_config(config, "tokenizer") or {}
|
| 222 |
+
tokenizer_path = tokenizer_config.get("path") or tokenizer_config.get("tokenizer_path")
|
| 223 |
+
tokenizer_type = tokenizer_config.get("type", "bpe")
|
| 224 |
+
auto_train = bool(tokenizer_config.get("auto_train", False))
|
| 225 |
+
if tokenizer_type in {"bpe", "sentencepiece"} and not tokenizer_path:
|
| 226 |
+
train_texts = _get_tokenizer_train_texts(config, allow_auto=auto_train)
|
| 227 |
+
if train_texts is None:
|
| 228 |
+
raise ValueError(
|
| 229 |
+
"BPE tokenizer requires tokenizer.path or a local custom dataset with train texts. "
|
| 230 |
+
"Automatic download from WMT/OPUS is disabled by default. "
|
| 231 |
+
"Set tokenizer.auto_train=true to enable it, or provide tokenizer.path/pretrained tokenizer."
|
| 232 |
+
)
|
| 233 |
+
tokenizer = build_tokenizer(tokenizer_config, train_texts=train_texts)
|
| 234 |
+
else:
|
| 235 |
+
tokenizer = build_tokenizer(tokenizer_config)
|
| 236 |
+
|
| 237 |
+
model = TransformerTranslationModel(
|
| 238 |
+
src_vocab_size=tokenizer.vocab_size,
|
| 239 |
+
tgt_vocab_size=tokenizer.vocab_size,
|
| 240 |
+
d_model=_get_config(config, "model", "transformer", "d_model"),
|
| 241 |
+
nhead=_get_config(config, "model", "transformer", "nhead"),
|
| 242 |
+
num_encoder_layers=_get_config(config, "model", "transformer", "num_encoder_layers"),
|
| 243 |
+
num_decoder_layers=_get_config(config, "model", "transformer", "num_decoder_layers"),
|
| 244 |
+
dim_feedforward=_get_config(config, "model", "transformer", "dim_feedforward"),
|
| 245 |
+
dropout=_get_config(config, "model", "transformer", "dropout"),
|
| 246 |
+
activation=_get_config(config, "model", "transformer", "activation"),
|
| 247 |
+
max_seq_len=_get_config(config, "model", "transformer", "max_seq_len"),
|
| 248 |
+
use_flash_attention=_get_config(config, "model", "transformer", "use_flash_attention"),
|
| 249 |
+
use_rotary_embedding=_get_config(config, "model", "transformer", "use_rotary_embedding"),
|
| 250 |
+
pre_norm=_get_config(config, "model", "transformer", "pre_norm"),
|
| 251 |
+
pad_id=tokenizer.pad_token_id,
|
| 252 |
+
)
|
| 253 |
+
return model.to(device), tokenizer
|
| 254 |
+
|
| 255 |
+
model, hf_tokenizer = load_pretrained_model(
|
| 256 |
+
config.model.pretrained.model_name,
|
| 257 |
+
config.model.pretrained.src_lang,
|
| 258 |
+
config.model.pretrained.tgt_lang,
|
| 259 |
+
device=str(device),
|
| 260 |
+
)
|
| 261 |
+
tokenizer = TokenizerWrapper(
|
| 262 |
+
hf_tokenizer,
|
| 263 |
+
pad_token=getattr(hf_tokenizer, "pad_token", "<pad>"),
|
| 264 |
+
unk_token=getattr(hf_tokenizer, "unk_token", "<unk>"),
|
| 265 |
+
bos_token=getattr(hf_tokenizer, "bos_token", "<s>"),
|
| 266 |
+
eos_token=getattr(hf_tokenizer, "eos_token", "</s>"),
|
| 267 |
+
)
|
| 268 |
+
return model, tokenizer
|
| 269 |
+
|
| 270 |
+
|
| 271 |
+
def load_checkpoint(model, checkpoint_path, device):
|
| 272 |
+
checkpoint = torch.load(checkpoint_path, map_location=device)
|
| 273 |
+
if isinstance(checkpoint, dict):
|
| 274 |
+
if "model_state_dict" in checkpoint:
|
| 275 |
+
model.load_state_dict(checkpoint["model_state_dict"])
|
| 276 |
+
elif "state_dict" in checkpoint:
|
| 277 |
+
model.load_state_dict(checkpoint["state_dict"])
|
| 278 |
+
else:
|
| 279 |
+
try:
|
| 280 |
+
model.load_state_dict(checkpoint)
|
| 281 |
+
except Exception as exc:
|
| 282 |
+
raise ValueError("Checkpoint does not contain a valid model state dict") from exc
|
| 283 |
+
else:
|
| 284 |
+
raise ValueError("Unsupported checkpoint format")
|
| 285 |
+
return model
|
| 286 |
+
|
| 287 |
+
|
| 288 |
+
def main():
|
| 289 |
+
args, cli_overrides = parse_args()
|
| 290 |
+
|
| 291 |
+
print("=" * 60)
|
| 292 |
+
print(" EasyTranslate - Translation")
|
| 293 |
+
print("=" * 60)
|
| 294 |
+
|
| 295 |
+
if args.web and not args.streamlit_app:
|
| 296 |
+
launch_streamlit_process(args)
|
| 297 |
+
return
|
| 298 |
+
|
| 299 |
+
if args.streamlit_app:
|
| 300 |
+
launch_streamlit_app(args)
|
| 301 |
+
return
|
| 302 |
+
|
| 303 |
+
config = _load_config(args.config, cli_overrides)
|
| 304 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 305 |
+
model, tokenizer = build_model_and_tokenizer(config, device)
|
| 306 |
+
model = load_checkpoint(model, args.checkpoint, device)
|
| 307 |
+
evaluator = Evaluator(model, tokenizer, config)
|
| 308 |
+
|
| 309 |
+
if args.input and args.output:
|
| 310 |
+
translate_file(evaluator, args.input, args.output)
|
| 311 |
+
else:
|
| 312 |
+
interactive_translate(evaluator)
|
| 313 |
+
|
| 314 |
+
|
| 315 |
+
if __name__ == "__main__":
|
| 316 |
+
main()
|
src/easytranslate/evaluation/__init__.py
CHANGED
|
@@ -1,16 +1,16 @@
|
|
| 1 |
-
"""评估与推理模块 (Person D 负责)"""
|
| 2 |
-
|
| 3 |
-
from easytranslate.evaluation.metrics import compute_bleu, compute_comet, compute_chrf, compute_all_metrics
|
| 4 |
-
from easytranslate.evaluation.decoding import greedy_decode, beam_search_decode, sample_decode
|
| 5 |
-
from easytranslate.evaluation.evaluator import Evaluator
|
| 6 |
-
|
| 7 |
-
__all__ = [
|
| 8 |
-
"compute_bleu",
|
| 9 |
-
"compute_comet",
|
| 10 |
-
"compute_chrf",
|
| 11 |
-
"compute_all_metrics",
|
| 12 |
-
"greedy_decode",
|
| 13 |
-
"beam_search_decode",
|
| 14 |
-
"sample_decode",
|
| 15 |
-
"Evaluator",
|
| 16 |
-
]
|
|
|
|
| 1 |
+
"""评估与推理模块 (Person D 负责)"""
|
| 2 |
+
|
| 3 |
+
from easytranslate.evaluation.metrics import compute_bleu, compute_comet, compute_chrf, compute_all_metrics
|
| 4 |
+
from easytranslate.evaluation.decoding import greedy_decode, beam_search_decode, sample_decode
|
| 5 |
+
from easytranslate.evaluation.evaluator import Evaluator
|
| 6 |
+
|
| 7 |
+
__all__ = [
|
| 8 |
+
"compute_bleu",
|
| 9 |
+
"compute_comet",
|
| 10 |
+
"compute_chrf",
|
| 11 |
+
"compute_all_metrics",
|
| 12 |
+
"greedy_decode",
|
| 13 |
+
"beam_search_decode",
|
| 14 |
+
"sample_decode",
|
| 15 |
+
"Evaluator",
|
| 16 |
+
]
|
src/easytranslate/evaluation/decoding.py
CHANGED
|
@@ -1,129 +1,219 @@
|
|
| 1 |
-
"""
|
| 2 |
-
解码策略模块 — Person D 负责实现
|
| 3 |
-
|
| 4 |
-
功能要求:
|
| 5 |
-
1. greedy_decode: 贪心解码
|
| 6 |
-
2. beam_search_decode: 束搜索解码
|
| 7 |
-
3. sample_decode: 采样解码 (temperature, top-k, top-p)
|
| 8 |
-
|
| 9 |
-
技术要点:
|
| 10 |
-
- Beam Search 是翻译任务最常用的解码策略
|
| 11 |
-
- 需要高效处理批量解码
|
| 12 |
-
- 支持长度惩罚 (length penalty) 和重复惩罚 (no_repeat_ngram)
|
| 13 |
-
- 对于预训练模型,可以直接使用 model.generate()
|
| 14 |
-
"""
|
| 15 |
-
|
| 16 |
-
from __future__ import annotations
|
| 17 |
-
|
| 18 |
-
import logging
|
| 19 |
-
from typing import Optional
|
| 20 |
-
|
| 21 |
-
import torch
|
| 22 |
-
import torch.nn as nn
|
| 23 |
-
import torch.nn.functional as F
|
| 24 |
-
|
| 25 |
-
logger = logging.getLogger(__name__)
|
| 26 |
-
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
|
| 31 |
-
|
| 32 |
-
|
| 33 |
-
|
| 34 |
-
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
|
| 46 |
-
|
| 47 |
-
|
| 48 |
-
|
| 49 |
-
|
| 50 |
-
|
| 51 |
-
|
| 52 |
-
|
| 53 |
-
|
| 54 |
-
|
| 55 |
-
|
| 56 |
-
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
|
| 61 |
-
|
| 62 |
-
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
|
| 67 |
-
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
|
| 71 |
-
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
|
| 75 |
-
|
| 76 |
-
|
| 77 |
-
|
| 78 |
-
|
| 79 |
-
|
| 80 |
-
|
| 81 |
-
|
| 82 |
-
|
| 83 |
-
|
| 84 |
-
|
| 85 |
-
|
| 86 |
-
|
| 87 |
-
|
| 88 |
-
|
| 89 |
-
|
| 90 |
-
|
| 91 |
-
|
| 92 |
-
|
| 93 |
-
|
| 94 |
-
|
| 95 |
-
|
| 96 |
-
|
| 97 |
-
|
| 98 |
-
|
| 99 |
-
|
| 100 |
-
|
| 101 |
-
|
| 102 |
-
|
| 103 |
-
|
| 104 |
-
|
| 105 |
-
|
| 106 |
-
|
| 107 |
-
|
| 108 |
-
1.
|
| 109 |
-
|
| 110 |
-
|
| 111 |
-
|
| 112 |
-
|
| 113 |
-
""
|
| 114 |
-
|
| 115 |
-
|
| 116 |
-
|
| 117 |
-
|
| 118 |
-
|
| 119 |
-
|
| 120 |
-
|
| 121 |
-
|
| 122 |
-
|
| 123 |
-
|
| 124 |
-
|
| 125 |
-
|
| 126 |
-
|
| 127 |
-
|
| 128 |
-
|
| 129 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
解码策略模块 — Person D 负责实现
|
| 3 |
+
|
| 4 |
+
功能要求:
|
| 5 |
+
1. greedy_decode: 贪心解码
|
| 6 |
+
2. beam_search_decode: 束搜索解码
|
| 7 |
+
3. sample_decode: 采样解码 (temperature, top-k, top-p)
|
| 8 |
+
|
| 9 |
+
技术要点:
|
| 10 |
+
- Beam Search 是翻译任务最常用的解码策略
|
| 11 |
+
- 需要高效处理批量解码
|
| 12 |
+
- 支持长度惩罚 (length penalty) 和重复惩罚 (no_repeat_ngram)
|
| 13 |
+
- 对于预训练模型,可以直接使用 model.generate()
|
| 14 |
+
"""
|
| 15 |
+
|
| 16 |
+
from __future__ import annotations
|
| 17 |
+
|
| 18 |
+
import logging
|
| 19 |
+
from typing import Optional
|
| 20 |
+
|
| 21 |
+
import torch
|
| 22 |
+
import torch.nn as nn
|
| 23 |
+
import torch.nn.functional as F
|
| 24 |
+
|
| 25 |
+
logger = logging.getLogger(__name__)
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def _apply_no_repeat_ngram(
|
| 29 |
+
logits: torch.Tensor,
|
| 30 |
+
generated_tokens: torch.Tensor,
|
| 31 |
+
ngram_size: int,
|
| 32 |
+
) -> torch.Tensor:
|
| 33 |
+
"""
|
| 34 |
+
防止生成重复的 n-gram。
|
| 35 |
+
"""
|
| 36 |
+
if ngram_size <= 1:
|
| 37 |
+
return logits
|
| 38 |
+
|
| 39 |
+
batch_size, vocab_size = logits.size()
|
| 40 |
+
for batch_idx in range(batch_size):
|
| 41 |
+
tokens = generated_tokens[batch_idx].tolist()
|
| 42 |
+
if len(tokens) < ngram_size - 1:
|
| 43 |
+
continue
|
| 44 |
+
|
| 45 |
+
banned_tokens: set[int] = set()
|
| 46 |
+
ngram_map: dict[tuple[int, ...], set[int]] = {}
|
| 47 |
+
for i in range(len(tokens) - ngram_size + 1):
|
| 48 |
+
prefix = tuple(tokens[i : i + ngram_size - 1])
|
| 49 |
+
next_token = tokens[i + ngram_size - 1]
|
| 50 |
+
ngram_map.setdefault(prefix, set()).add(next_token)
|
| 51 |
+
|
| 52 |
+
prefix = tuple(tokens[-(ngram_size - 1) :])
|
| 53 |
+
if prefix in ngram_map:
|
| 54 |
+
banned_tokens = ngram_map[prefix]
|
| 55 |
+
logits[batch_idx, list(banned_tokens)] = float("-inf")
|
| 56 |
+
|
| 57 |
+
return logits
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
@torch.no_grad()
|
| 61 |
+
def greedy_decode(
|
| 62 |
+
model: nn.Module,
|
| 63 |
+
src_ids: torch.Tensor,
|
| 64 |
+
src_padding_mask: torch.BoolTensor,
|
| 65 |
+
bos_id: int,
|
| 66 |
+
eos_id: int,
|
| 67 |
+
max_len: int = 256,
|
| 68 |
+
) -> torch.Tensor:
|
| 69 |
+
"""
|
| 70 |
+
贪心解码。
|
| 71 |
+
"""
|
| 72 |
+
encoder_output = model.encode(src_ids, src_padding_mask)
|
| 73 |
+
batch_size = src_ids.size(0)
|
| 74 |
+
device = src_ids.device
|
| 75 |
+
generated = torch.full((batch_size, 1), bos_id, dtype=torch.long, device=device)
|
| 76 |
+
finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
|
| 77 |
+
|
| 78 |
+
for _ in range(max_len):
|
| 79 |
+
logits = model.decode_step(generated, encoder_output, src_padding_mask)
|
| 80 |
+
next_token = logits.argmax(dim=-1, keepdim=True)
|
| 81 |
+
generated = torch.cat([generated, next_token], dim=1)
|
| 82 |
+
finished = finished | next_token.squeeze(-1).eq(eos_id)
|
| 83 |
+
if finished.all():
|
| 84 |
+
break
|
| 85 |
+
|
| 86 |
+
return generated
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
@torch.no_grad()
|
| 90 |
+
def beam_search_decode(
|
| 91 |
+
model: nn.Module,
|
| 92 |
+
src_ids: torch.Tensor,
|
| 93 |
+
src_padding_mask: torch.BoolTensor,
|
| 94 |
+
bos_id: int,
|
| 95 |
+
eos_id: int,
|
| 96 |
+
beam_size: int = 5,
|
| 97 |
+
max_len: int = 256,
|
| 98 |
+
length_penalty: float = 1.0,
|
| 99 |
+
no_repeat_ngram_size: int = 0,
|
| 100 |
+
) -> torch.Tensor:
|
| 101 |
+
"""
|
| 102 |
+
束搜索解码。
|
| 103 |
+
"""
|
| 104 |
+
batch_size, seq_len = src_ids.size()
|
| 105 |
+
device = src_ids.device
|
| 106 |
+
|
| 107 |
+
encoder_output = model.encode(src_ids, src_padding_mask)
|
| 108 |
+
encoder_output = encoder_output.unsqueeze(1).expand(batch_size, beam_size, -1, -1)
|
| 109 |
+
encoder_output = encoder_output.reshape(batch_size * beam_size, seq_len, -1)
|
| 110 |
+
src_padding_mask = src_padding_mask.unsqueeze(1).expand(batch_size, beam_size, seq_len)
|
| 111 |
+
src_padding_mask = src_padding_mask.reshape(batch_size * beam_size, seq_len)
|
| 112 |
+
|
| 113 |
+
beam_scores = torch.full((batch_size, beam_size), float("-inf"), device=device)
|
| 114 |
+
beam_scores[:, 0] = 0.0
|
| 115 |
+
generated = torch.full((batch_size, beam_size, 1), bos_id, dtype=torch.long, device=device)
|
| 116 |
+
finished = torch.zeros((batch_size, beam_size), dtype=torch.bool, device=device)
|
| 117 |
+
|
| 118 |
+
for _ in range(max_len):
|
| 119 |
+
flat_generated = generated.view(batch_size * beam_size, -1)
|
| 120 |
+
logits = model.decode_step(flat_generated, encoder_output, src_padding_mask)
|
| 121 |
+
log_probs = F.log_softmax(logits, dim=-1)
|
| 122 |
+
|
| 123 |
+
if no_repeat_ngram_size > 0:
|
| 124 |
+
log_probs = _apply_no_repeat_ngram(log_probs, flat_generated, no_repeat_ngram_size)
|
| 125 |
+
|
| 126 |
+
finished_flat = finished.view(batch_size * beam_size)
|
| 127 |
+
if finished_flat.any():
|
| 128 |
+
log_probs[finished_flat] = float("-inf")
|
| 129 |
+
log_probs[finished_flat, eos_id] = 0.0
|
| 130 |
+
|
| 131 |
+
vocab_size = log_probs.size(-1)
|
| 132 |
+
scores = beam_scores.unsqueeze(-1) + log_probs.view(batch_size, beam_size, vocab_size)
|
| 133 |
+
scores_flat = scores.view(batch_size, -1)
|
| 134 |
+
topk_scores, topk_indices = scores_flat.topk(beam_size, dim=-1)
|
| 135 |
+
|
| 136 |
+
beam_indices = topk_indices // vocab_size
|
| 137 |
+
token_indices = topk_indices % vocab_size
|
| 138 |
+
|
| 139 |
+
next_generated = []
|
| 140 |
+
next_finished = []
|
| 141 |
+
for batch_idx in range(batch_size):
|
| 142 |
+
selected_beams = beam_indices[batch_idx]
|
| 143 |
+
selected_tokens = token_indices[batch_idx]
|
| 144 |
+
next_seq = generated[batch_idx, selected_beams]
|
| 145 |
+
next_seq = torch.cat([next_seq, selected_tokens.unsqueeze(-1)], dim=-1)
|
| 146 |
+
next_generated.append(next_seq)
|
| 147 |
+
next_finished.append(
|
| 148 |
+
finished[batch_idx, selected_beams]
|
| 149 |
+
| selected_tokens.eq(eos_id)
|
| 150 |
+
)
|
| 151 |
+
|
| 152 |
+
generated = torch.stack(next_generated, dim=0)
|
| 153 |
+
finished = torch.stack(next_finished, dim=0)
|
| 154 |
+
beam_scores = topk_scores
|
| 155 |
+
|
| 156 |
+
if finished.all():
|
| 157 |
+
break
|
| 158 |
+
|
| 159 |
+
length = generated.size(1)
|
| 160 |
+
penalty = float(length) ** float(length_penalty)
|
| 161 |
+
final_scores = beam_scores / penalty
|
| 162 |
+
best_indices = final_scores.argmax(dim=-1)
|
| 163 |
+
output = generated[torch.arange(batch_size, device=device), best_indices]
|
| 164 |
+
return output
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
@torch.no_grad()
|
| 168 |
+
def sample_decode(
|
| 169 |
+
model: nn.Module,
|
| 170 |
+
src_ids: torch.Tensor,
|
| 171 |
+
src_padding_mask: torch.BoolTensor,
|
| 172 |
+
bos_id: int,
|
| 173 |
+
eos_id: int,
|
| 174 |
+
max_len: int = 256,
|
| 175 |
+
temperature: float = 1.0,
|
| 176 |
+
top_k: int = 0,
|
| 177 |
+
top_p: float = 1.0,
|
| 178 |
+
) -> torch.Tensor:
|
| 179 |
+
"""
|
| 180 |
+
采样解码 (支持 temperature, top-k, top-p/nucleus sampling)。
|
| 181 |
+
"""
|
| 182 |
+
encoder_output = model.encode(src_ids, src_padding_mask)
|
| 183 |
+
batch_size = src_ids.size(0)
|
| 184 |
+
device = src_ids.device
|
| 185 |
+
generated = torch.full((batch_size, 1), bos_id, dtype=torch.long, device=device)
|
| 186 |
+
finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
|
| 187 |
+
|
| 188 |
+
for _ in range(max_len):
|
| 189 |
+
logits = model.decode_step(generated, encoder_output, src_padding_mask)
|
| 190 |
+
logits = logits / max(temperature, 1e-8)
|
| 191 |
+
|
| 192 |
+
if top_k > 0:
|
| 193 |
+
top_k = min(top_k, logits.size(-1))
|
| 194 |
+
values, indices = torch.topk(logits, top_k, dim=-1)
|
| 195 |
+
mask = torch.full_like(logits, float("-inf"))
|
| 196 |
+
logits = mask.scatter(-1, indices, values)
|
| 197 |
+
|
| 198 |
+
if 0.0 < top_p < 1.0:
|
| 199 |
+
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
|
| 200 |
+
probs = F.softmax(sorted_logits, dim=-1)
|
| 201 |
+
cumulative_probs = torch.cumsum(probs, dim=-1)
|
| 202 |
+
cutoff = cumulative_probs > top_p
|
| 203 |
+
cutoff[:, 1:] = cutoff[:, :-1].clone()
|
| 204 |
+
cutoff[:, 0] = False
|
| 205 |
+
sorted_logits[cutoff] = float("-inf")
|
| 206 |
+
logits = torch.zeros_like(logits).scatter(-1, sorted_indices, sorted_logits)
|
| 207 |
+
|
| 208 |
+
probs = F.softmax(logits, dim=-1)
|
| 209 |
+
next_token = torch.multinomial(probs, num_samples=1)
|
| 210 |
+
next_token = next_token.clamp(min=0)
|
| 211 |
+
|
| 212 |
+
next_token = torch.where(finished.unsqueeze(-1), torch.full_like(next_token, eos_id), next_token)
|
| 213 |
+
generated = torch.cat([generated, next_token], dim=1)
|
| 214 |
+
finished = finished | next_token.squeeze(-1).eq(eos_id)
|
| 215 |
+
|
| 216 |
+
if finished.all():
|
| 217 |
+
break
|
| 218 |
+
|
| 219 |
+
return generated
|
src/easytranslate/evaluation/evaluator.py
CHANGED
|
@@ -1,76 +1,208 @@
|
|
| 1 |
-
"""
|
| 2 |
-
评估器模块 — Person D 负责实现
|
| 3 |
-
|
| 4 |
-
功能要求:
|
| 5 |
-
将解码和评估指标整合为统一的评估接口。
|
| 6 |
-
|
| 7 |
-
使用方法:
|
| 8 |
-
evaluator = Evaluator(model, tokenizer, config)
|
| 9 |
-
results = evaluator.evaluate(test_loader)
|
| 10 |
-
"""
|
| 11 |
-
|
| 12 |
-
from __future__ import annotations
|
| 13 |
-
|
| 14 |
-
import logging
|
| 15 |
-
from typing import Optional
|
| 16 |
-
|
| 17 |
-
import torch
|
| 18 |
-
import torch.nn as nn
|
| 19 |
-
from torch.utils.
|
| 20 |
-
from
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
from easytranslate.evaluation.
|
| 24 |
-
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
|
| 31 |
-
|
| 32 |
-
|
| 33 |
-
|
| 34 |
-
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
|
| 46 |
-
|
| 47 |
-
|
| 48 |
-
|
| 49 |
-
|
| 50 |
-
|
| 51 |
-
|
| 52 |
-
|
| 53 |
-
|
| 54 |
-
|
| 55 |
-
|
| 56 |
-
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
"""
|
| 60 |
-
|
| 61 |
-
|
| 62 |
-
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
|
| 67 |
-
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
|
| 71 |
-
"""
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
|
| 75 |
-
|
| 76 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
评估器模块 — Person D 负责实现
|
| 3 |
+
|
| 4 |
+
功能要求:
|
| 5 |
+
将解码和评估指标整合为统一的评估接口。
|
| 6 |
+
|
| 7 |
+
使用方法:
|
| 8 |
+
evaluator = Evaluator(model, tokenizer, config)
|
| 9 |
+
results = evaluator.evaluate(test_loader)
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
from __future__ import annotations
|
| 13 |
+
|
| 14 |
+
import logging
|
| 15 |
+
from typing import Optional
|
| 16 |
+
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
from torch.nn.utils.rnn import pad_sequence
|
| 20 |
+
from torch.utils.data import DataLoader
|
| 21 |
+
from tqdm import tqdm
|
| 22 |
+
|
| 23 |
+
from easytranslate.evaluation.metrics import compute_all_metrics
|
| 24 |
+
from easytranslate.evaluation.decoding import greedy_decode, beam_search_decode, sample_decode
|
| 25 |
+
|
| 26 |
+
logger = logging.getLogger(__name__)
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def _get_config_value(config, key_path, default=None):
|
| 30 |
+
if config is None:
|
| 31 |
+
return default
|
| 32 |
+
if isinstance(config, dict):
|
| 33 |
+
value = config
|
| 34 |
+
for key in key_path:
|
| 35 |
+
value = value.get(key, default)
|
| 36 |
+
if value is default:
|
| 37 |
+
break
|
| 38 |
+
return value
|
| 39 |
+
value = config
|
| 40 |
+
for key in key_path:
|
| 41 |
+
value = getattr(value, key, default)
|
| 42 |
+
if value is default:
|
| 43 |
+
break
|
| 44 |
+
return value
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
class Evaluator:
|
| 48 |
+
"""
|
| 49 |
+
翻译模型评估器。
|
| 50 |
+
"""
|
| 51 |
+
|
| 52 |
+
def __init__(self, model: nn.Module, tokenizer, config: dict):
|
| 53 |
+
self.model = model
|
| 54 |
+
self.tokenizer = tokenizer
|
| 55 |
+
self.config = config or {}
|
| 56 |
+
|
| 57 |
+
self.evaluation_config = _get_config_value(self.config, ["evaluation"], {})
|
| 58 |
+
self.decoding_config = _get_config_value(self.evaluation_config, ["decoding"], {})
|
| 59 |
+
self.metrics = _get_config_value(self.evaluation_config, ["metrics"], ["bleu", "comet", "chrf", "ter"])
|
| 60 |
+
self.strategy = self.decoding_config.get("strategy", "beam_search") if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "strategy", "beam_search")
|
| 61 |
+
|
| 62 |
+
self.bos_id = self.tokenizer.bos_token_id
|
| 63 |
+
self.eos_id = self.tokenizer.eos_token_id
|
| 64 |
+
self.pad_id = self.tokenizer.pad_token_id
|
| 65 |
+
|
| 66 |
+
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)
|
| 67 |
+
self.beam_size = self.decoding_config.get("beam_size", 5) if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "beam_size", 5)
|
| 68 |
+
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)
|
| 69 |
+
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)
|
| 70 |
+
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)
|
| 71 |
+
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)
|
| 72 |
+
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)
|
| 73 |
+
|
| 74 |
+
self.use_generate = hasattr(self.model, "generate") and not (hasattr(self.model, "encode") and hasattr(self.model, "decode_step"))
|
| 75 |
+
|
| 76 |
+
def _decode(self, src_ids: torch.Tensor, src_padding_mask: torch.BoolTensor) -> torch.Tensor:
|
| 77 |
+
if self.use_generate:
|
| 78 |
+
generate_kwargs = {
|
| 79 |
+
"max_length": self.max_decode_len,
|
| 80 |
+
"early_stopping": True,
|
| 81 |
+
}
|
| 82 |
+
if self.strategy == "beam_search":
|
| 83 |
+
generate_kwargs.update(
|
| 84 |
+
{
|
| 85 |
+
"num_beams": self.beam_size,
|
| 86 |
+
"length_penalty": self.length_penalty,
|
| 87 |
+
"no_repeat_ngram_size": self.no_repeat_ngram_size,
|
| 88 |
+
}
|
| 89 |
+
)
|
| 90 |
+
elif self.strategy == "sampling":
|
| 91 |
+
generate_kwargs.update(
|
| 92 |
+
{
|
| 93 |
+
"do_sample": True,
|
| 94 |
+
"temperature": self.temperature,
|
| 95 |
+
"top_k": self.top_k,
|
| 96 |
+
"top_p": self.top_p,
|
| 97 |
+
"num_beams": 1,
|
| 98 |
+
}
|
| 99 |
+
)
|
| 100 |
+
else:
|
| 101 |
+
generate_kwargs.update({"num_beams": 1})
|
| 102 |
+
|
| 103 |
+
attention_mask = (~src_padding_mask).long()
|
| 104 |
+
return self.model.generate(input_ids=src_ids, attention_mask=attention_mask, **generate_kwargs)
|
| 105 |
+
|
| 106 |
+
if self.strategy == "beam_search":
|
| 107 |
+
return beam_search_decode(
|
| 108 |
+
self.model,
|
| 109 |
+
src_ids,
|
| 110 |
+
src_padding_mask,
|
| 111 |
+
self.bos_id,
|
| 112 |
+
self.eos_id,
|
| 113 |
+
beam_size=self.beam_size,
|
| 114 |
+
max_len=self.max_decode_len,
|
| 115 |
+
length_penalty=self.length_penalty,
|
| 116 |
+
no_repeat_ngram_size=self.no_repeat_ngram_size,
|
| 117 |
+
)
|
| 118 |
+
if self.strategy == "sampling":
|
| 119 |
+
return sample_decode(
|
| 120 |
+
self.model,
|
| 121 |
+
src_ids,
|
| 122 |
+
src_padding_mask,
|
| 123 |
+
self.bos_id,
|
| 124 |
+
self.eos_id,
|
| 125 |
+
max_len=self.max_decode_len,
|
| 126 |
+
temperature=self.temperature,
|
| 127 |
+
top_k=self.top_k,
|
| 128 |
+
top_p=self.top_p,
|
| 129 |
+
)
|
| 130 |
+
return greedy_decode(
|
| 131 |
+
self.model,
|
| 132 |
+
src_ids,
|
| 133 |
+
src_padding_mask,
|
| 134 |
+
self.bos_id,
|
| 135 |
+
self.eos_id,
|
| 136 |
+
max_len=self.max_decode_len,
|
| 137 |
+
)
|
| 138 |
+
|
| 139 |
+
def evaluate(
|
| 140 |
+
self,
|
| 141 |
+
dataloader: DataLoader,
|
| 142 |
+
src_texts: Optional[list[str]] = None,
|
| 143 |
+
ref_texts: Optional[list[str]] = None,
|
| 144 |
+
) -> dict:
|
| 145 |
+
self.model.eval()
|
| 146 |
+
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")
|
| 147 |
+
|
| 148 |
+
if src_texts is None and hasattr(dataloader, "dataset") and hasattr(dataloader.dataset, "src_texts"):
|
| 149 |
+
src_texts = list(dataloader.dataset.src_texts)
|
| 150 |
+
if ref_texts is None and hasattr(dataloader, "dataset") and hasattr(dataloader.dataset, "tgt_texts"):
|
| 151 |
+
ref_texts = list(dataloader.dataset.tgt_texts)
|
| 152 |
+
|
| 153 |
+
if src_texts is None or ref_texts is None:
|
| 154 |
+
raise ValueError("Source texts and reference texts must be provided for evaluation.")
|
| 155 |
+
|
| 156 |
+
hypotheses: list[str] = []
|
| 157 |
+
sources: list[str] = []
|
| 158 |
+
references: list[str] = []
|
| 159 |
+
|
| 160 |
+
for batch_idx, batch in enumerate(tqdm(dataloader, desc="Evaluating", unit="batch")):
|
| 161 |
+
src_ids = batch["src_ids"].to(device)
|
| 162 |
+
src_padding_mask = batch.get("src_padding_mask")
|
| 163 |
+
if src_padding_mask is None:
|
| 164 |
+
src_padding_mask = src_ids.eq(self.pad_id)
|
| 165 |
+
else:
|
| 166 |
+
src_padding_mask = src_padding_mask.to(device)
|
| 167 |
+
|
| 168 |
+
output_ids = self._decode(src_ids, src_padding_mask)
|
| 169 |
+
if isinstance(output_ids, torch.Tensor):
|
| 170 |
+
output_ids = output_ids.cpu()
|
| 171 |
+
|
| 172 |
+
for sample_idx in range(output_ids.size(0)):
|
| 173 |
+
decoded = self.tokenizer.decode(output_ids[sample_idx].tolist(), skip_special_tokens=True)
|
| 174 |
+
hypotheses.append(decoded)
|
| 175 |
+
|
| 176 |
+
sources = src_texts
|
| 177 |
+
references = ref_texts
|
| 178 |
+
results = compute_all_metrics(sources, hypotheses, references, metrics=self.metrics)
|
| 179 |
+
return results
|
| 180 |
+
|
| 181 |
+
def translate(self, texts: list[str]) -> list[str]:
|
| 182 |
+
self.model.eval()
|
| 183 |
+
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")
|
| 184 |
+
|
| 185 |
+
input_ids = []
|
| 186 |
+
for text in texts:
|
| 187 |
+
src_tokens = self.tokenizer.encode(
|
| 188 |
+
text,
|
| 189 |
+
add_special_tokens=True,
|
| 190 |
+
max_length=_get_config_value(self.config, ["data", "preprocessing", "max_src_len"], 256),
|
| 191 |
+
)
|
| 192 |
+
input_ids.append(torch.tensor(src_tokens, dtype=torch.long, device=device))
|
| 193 |
+
|
| 194 |
+
src_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.pad_id)
|
| 195 |
+
src_padding_mask = src_ids.eq(self.pad_id)
|
| 196 |
+
|
| 197 |
+
output_ids = self._decode(src_ids, src_padding_mask)
|
| 198 |
+
if isinstance(output_ids, torch.Tensor):
|
| 199 |
+
output_ids = output_ids.cpu()
|
| 200 |
+
|
| 201 |
+
translations: list[str] = []
|
| 202 |
+
for sample_idx in range(output_ids.size(0)):
|
| 203 |
+
translations.append(self.tokenizer.decode(output_ids[sample_idx].tolist(), skip_special_tokens=True))
|
| 204 |
+
|
| 205 |
+
return translations
|
| 206 |
+
|
| 207 |
+
def translate_single(self, text: str) -> str:
|
| 208 |
+
return self.translate([text])[0]
|
src/easytranslate/evaluation/metrics.py
CHANGED
|
@@ -1,109 +1,169 @@
|
|
| 1 |
-
"""
|
| 2 |
-
评估指标模块 — Person D 负责实现
|
| 3 |
-
|
| 4 |
-
功能要求:
|
| 5 |
-
1. compute_bleu: 计算 SacreBLEU 分数
|
| 6 |
-
2. compute_comet: 计算 COMET 分数 (神经网络评估指标)
|
| 7 |
-
3. compute_chrf: 计算 chrF++ 分数
|
| 8 |
-
4. compute_ter: 计算 TER (Translation Edit Rate)
|
| 9 |
-
5. compute_all_metrics: 计算所有指标
|
| 10 |
-
|
| 11 |
-
技术要点:
|
| 12 |
-
- SacreBLEU: 标准化的 BLEU 实现,结果可复现
|
| 13 |
-
- COMET: 基于预训练语言模型的评估指标,与人类评价相关性最高
|
| 14 |
-
- chrF++: 基于字符 n-gram 的 F-score,对中文尤其有用
|
| 15 |
-
- TER: 编辑距离,衡量翻译后编辑量
|
| 16 |
-
"""
|
| 17 |
-
|
| 18 |
-
from __future__ import annotations
|
| 19 |
-
|
| 20 |
-
import logging
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
|
| 31 |
-
|
| 32 |
-
|
| 33 |
-
|
| 34 |
-
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
|
| 46 |
-
|
| 47 |
-
references:
|
| 48 |
-
|
| 49 |
-
|
| 50 |
-
|
| 51 |
-
|
| 52 |
-
|
| 53 |
-
|
| 54 |
-
|
| 55 |
-
|
| 56 |
-
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
|
| 61 |
-
|
| 62 |
-
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
|
| 67 |
-
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
|
| 71 |
-
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
|
| 75 |
-
|
| 76 |
-
|
| 77 |
-
|
| 78 |
-
|
| 79 |
-
|
| 80 |
-
|
| 81 |
-
|
| 82 |
-
references:
|
| 83 |
-
|
| 84 |
-
|
| 85 |
-
|
| 86 |
-
|
| 87 |
-
|
| 88 |
-
|
| 89 |
-
|
| 90 |
-
|
| 91 |
-
|
| 92 |
-
|
| 93 |
-
|
| 94 |
-
|
| 95 |
-
|
| 96 |
-
|
| 97 |
-
|
| 98 |
-
|
| 99 |
-
|
| 100 |
-
|
| 101 |
-
|
| 102 |
-
|
| 103 |
-
|
| 104 |
-
|
| 105 |
-
|
| 106 |
-
|
| 107 |
-
|
| 108 |
-
|
| 109 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
评估指标模块 — Person D 负责实现
|
| 3 |
+
|
| 4 |
+
功能要求:
|
| 5 |
+
1. compute_bleu: 计算 SacreBLEU 分数
|
| 6 |
+
2. compute_comet: 计算 COMET 分数 (神经网络评估指标)
|
| 7 |
+
3. compute_chrf: 计算 chrF++ 分数
|
| 8 |
+
4. compute_ter: 计算 TER (Translation Edit Rate)
|
| 9 |
+
5. compute_all_metrics: 计算所有指标
|
| 10 |
+
|
| 11 |
+
技术要点:
|
| 12 |
+
- SacreBLEU: 标准化的 BLEU 实现,结果可复现
|
| 13 |
+
- COMET: 基于预训练语言模型的评估指标,与人类评价相关性最高
|
| 14 |
+
- chrF++: 基于字符 n-gram 的 F-score,对中文尤其有用
|
| 15 |
+
- TER: 编辑距离,衡量翻译后编辑量
|
| 16 |
+
"""
|
| 17 |
+
|
| 18 |
+
from __future__ import annotations
|
| 19 |
+
|
| 20 |
+
import logging
|
| 21 |
+
import time
|
| 22 |
+
from typing import Optional
|
| 23 |
+
|
| 24 |
+
logger = logging.getLogger(__name__)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def _import_sacrebleu():
|
| 28 |
+
try:
|
| 29 |
+
import sacrebleu
|
| 30 |
+
except ImportError as exc:
|
| 31 |
+
logger.error("SacreBLEU is not installed: %s", exc)
|
| 32 |
+
raise ImportError(
|
| 33 |
+
"sacrebleu is required for BLEU/chrF/TER evaluation. "
|
| 34 |
+
"Install it with `python -m pip install sacrebleu`."
|
| 35 |
+
) from exc
|
| 36 |
+
return sacrebleu
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def compute_bleu(
|
| 40 |
+
hypotheses: list[str],
|
| 41 |
+
references: list[str],
|
| 42 |
+
tokenize: str = "zh",
|
| 43 |
+
) -> dict:
|
| 44 |
+
"""
|
| 45 |
+
计算 SacreBLEU 分数。
|
| 46 |
+
"""
|
| 47 |
+
if not hypotheses or not references:
|
| 48 |
+
return {
|
| 49 |
+
"bleu": 0.0,
|
| 50 |
+
"bleu_1": 0.0,
|
| 51 |
+
"bleu_2": 0.0,
|
| 52 |
+
"bleu_3": 0.0,
|
| 53 |
+
"bleu_4": 0.0,
|
| 54 |
+
"bp": 0.0,
|
| 55 |
+
}
|
| 56 |
+
|
| 57 |
+
sacrebleu = _import_sacrebleu()
|
| 58 |
+
references_list = [references]
|
| 59 |
+
bleu = sacrebleu.corpus_bleu(hypotheses, references_list, tokenize=tokenize)
|
| 60 |
+
precisions = [round(float(x), 4) for x in bleu.precisions]
|
| 61 |
+
return {
|
| 62 |
+
"bleu": round(float(bleu.score), 4),
|
| 63 |
+
"bleu_1": precisions[0] if len(precisions) > 0 else 0.0,
|
| 64 |
+
"bleu_2": precisions[1] if len(precisions) > 1 else 0.0,
|
| 65 |
+
"bleu_3": precisions[2] if len(precisions) > 2 else 0.0,
|
| 66 |
+
"bleu_4": precisions[3] if len(precisions) > 3 else 0.0,
|
| 67 |
+
"bp": round(float(getattr(bleu, "bp", 0.0)), 4),
|
| 68 |
+
}
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
def compute_comet(
|
| 72 |
+
sources: list[str],
|
| 73 |
+
hypotheses: list[str],
|
| 74 |
+
references: list[str],
|
| 75 |
+
model_name: str = "Unbabel/wmt22-comet-da",
|
| 76 |
+
batch_size: int = 16,
|
| 77 |
+
gpus: int = 1,
|
| 78 |
+
) -> dict:
|
| 79 |
+
"""
|
| 80 |
+
计算 COMET 分数。
|
| 81 |
+
"""
|
| 82 |
+
if not sources or not hypotheses or not references:
|
| 83 |
+
return {"comet": 0.0, "comet_scores": []}
|
| 84 |
+
|
| 85 |
+
try:
|
| 86 |
+
from comet import download_model, load_from_checkpoint
|
| 87 |
+
except ImportError as exc:
|
| 88 |
+
logger.error("COMET library is not installed: %s", exc)
|
| 89 |
+
raise
|
| 90 |
+
|
| 91 |
+
model_path = download_model(model_name)
|
| 92 |
+
model = load_from_checkpoint(model_path)
|
| 93 |
+
|
| 94 |
+
data = [
|
| 95 |
+
{"src": src, "mt": hyp, "ref": ref}
|
| 96 |
+
for src, hyp, ref in zip(sources, hypotheses, references)
|
| 97 |
+
]
|
| 98 |
+
prediction = model.predict(data, batch_size=batch_size, gpus=gpus)
|
| 99 |
+
|
| 100 |
+
if isinstance(prediction, dict):
|
| 101 |
+
scores = prediction.get("scores") or prediction.get("predictions") or []
|
| 102 |
+
else:
|
| 103 |
+
scores = list(prediction)
|
| 104 |
+
|
| 105 |
+
if scores is None:
|
| 106 |
+
scores = []
|
| 107 |
+
|
| 108 |
+
scores = [float(score) for score in scores]
|
| 109 |
+
system_score = float(sum(scores) / len(scores)) if scores else 0.0
|
| 110 |
+
return {"comet": round(system_score, 4), "comet_scores": scores}
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
def compute_chrf(
|
| 114 |
+
hypotheses: list[str],
|
| 115 |
+
references: list[str],
|
| 116 |
+
) -> dict:
|
| 117 |
+
"""
|
| 118 |
+
计算 chrF++ 分数。
|
| 119 |
+
"""
|
| 120 |
+
if not hypotheses or not references:
|
| 121 |
+
return {"chrf": 0.0}
|
| 122 |
+
|
| 123 |
+
sacrebleu = _import_sacrebleu()
|
| 124 |
+
chrf = sacrebleu.corpus_chrf(hypotheses, [references])
|
| 125 |
+
return {"chrf": round(float(chrf.score), 4)}
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
def compute_ter(
|
| 129 |
+
hypotheses: list[str],
|
| 130 |
+
references: list[str],
|
| 131 |
+
) -> dict:
|
| 132 |
+
"""
|
| 133 |
+
计算 TER 分数。
|
| 134 |
+
"""
|
| 135 |
+
if not hypotheses or not references:
|
| 136 |
+
return {"ter": 0.0}
|
| 137 |
+
|
| 138 |
+
sacrebleu = _import_sacrebleu()
|
| 139 |
+
ter = sacrebleu.corpus_ter(hypotheses, [references])
|
| 140 |
+
return {"ter": round(float(ter.score), 4)}
|
| 141 |
+
|
| 142 |
+
|
| 143 |
+
def compute_all_metrics(
|
| 144 |
+
sources: list[str],
|
| 145 |
+
hypotheses: list[str],
|
| 146 |
+
references: list[str],
|
| 147 |
+
metrics: list[str] = ["bleu", "comet", "chrf", "ter"],
|
| 148 |
+
) -> dict:
|
| 149 |
+
"""
|
| 150 |
+
计算所有指定的评估指标。
|
| 151 |
+
"""
|
| 152 |
+
results: dict[str, float | list[float]] = {}
|
| 153 |
+
|
| 154 |
+
for metric in metrics:
|
| 155 |
+
start_time = time.time()
|
| 156 |
+
if metric == "bleu":
|
| 157 |
+
results.update(compute_bleu(hypotheses, references))
|
| 158 |
+
elif metric == "comet":
|
| 159 |
+
results.update(compute_comet(sources, hypotheses, references))
|
| 160 |
+
elif metric == "chrf":
|
| 161 |
+
results.update(compute_chrf(hypotheses, references))
|
| 162 |
+
elif metric == "ter":
|
| 163 |
+
results.update(compute_ter(hypotheses, references))
|
| 164 |
+
else:
|
| 165 |
+
logger.warning("Unsupported evaluation metric: %s", metric)
|
| 166 |
+
elapsed = time.time() - start_time
|
| 167 |
+
logger.info("Computed %s in %.2f seconds", metric, elapsed)
|
| 168 |
+
|
| 169 |
+
return results
|