# 📊 MultiMetric-Eval

多指标翻译评测工具，一行代码计算 BLEU、chrF++、COMET、BLEURT，支持文本和语音双模式输入与评测。

[![PyPI version](https://badge.fury.io/py/multimetric-eval.svg)](https://badge.fury.io/py/multimetric-eval)
[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

---

## 🚀 安装

```bash
# 基础安装（BLEU + chrF++）
pip install multimetric-eval

# 按需安装可选依赖
pip install multimetric-eval[comet]     # + COMET 指标
pip install multimetric-eval[whisper]   # + 语音转文字
pip install multimetric-eval[all]       # 全部功能（不含 BLEURT）
```

### BLEURT 安装（PyTorch 版）

本工具使用 **PyTorch 原生版 BLEURT**（[bleurt-pytorch](https://github.com/lucadiliello/bleurt-pytorch)），与 COMET 共享 PyTorch 环境，**无需安装 TensorFlow，无环境冲突**。

```bash
# 安装 bleurt-pytorch
pip install git+https://github.com/lucadiliello/bleurt-pytorch.git
```

#### BLEURT 模型选择

| 模型 | HuggingFace ID | 大小 | 说明 |
|------|----------------|------|------|
| **BLEURT-20** | `lucadiliello/BLEURT-20` | ~3.6GB | 推荐，效果最好 |
| **BLEURT-20-D12** | `lucadiliello/BLEURT-20-D12` | ~1.1GB | 轻量版，适合资源有限场景 |

#### 使用方式

**方式 A：在线自动下载（推荐）**

无需手动下载，初始化时自动从 HuggingFace 下载并缓存：

```python
evaluator = ModelEvaluator(use_bleurt=True)
# 默认下载 lucadiliello/BLEURT-20

evaluator = ModelEvaluator(use_bleurt=True, bleurt_model="lucadiliello/BLEURT-20-D12")
# 指定使用轻量版
```

**方式 B：离线使用本地模型**

适用于无法联网的服务器：

```bash
# 1. 在联网机器上下载模型
python -c "
from huggingface_hub import snapshot_download
snapshot_download('lucadiliello/BLEURT-20', local_dir='./BLEURT-20', local_dir_use_symlinks=False)
"

# 2. 打包上传到离线机器
tar -czvf BLEURT-20.tar.gz BLEURT-20/

# 3. 在离线机器上解压后使用
```

```python
evaluator = ModelEvaluator(
    use_bleurt=True,
    bleurt_path="./BLEURT-20",  # 指定本地路径
)
```

> 📌 **优先级**：`bleurt_path`（本地路径）> `bleurt_model`（在线下载）> 默认模型（`lucadiliello/BLEURT-20`）

---

## 📖 快速开始

```python
from multimetric_eval import ModelEvaluator

# 初始化（首次会自动下载 COMET 模型）
evaluator = ModelEvaluator()

# 评测
results = evaluator.evaluate(
    hypothesis=["The cat sits on the mat."],
    reference=["The cat is sitting on the mat."],
    source=["猫坐在垫子上。"]
)

print(results)
# {'sacreBLEU': 45.23, 'chrF++': 62.15, 'COMET': 0.8523}
```

---

## 🎯 三种评测模式

本工具支持 **文本 (target_text)** 和 **语音 (target_speech)** 两种输入，可单独使用或同时使用：

| 模式 | 输入 | 输出指标 |
|------|------|----------|
| 纯文本 | `target_text` | `sacreBLEU`, `chrF++`, `COMET`, `BLEURT` |
| 纯语音 | `target_speech` | `sacreBLEU_ASR`, `chrF++_ASR`, `COMET_ASR`, `BLEURT_ASR` |
| 双模式 | 两者同时 | 以上全部指标 |

### 纯文本评测

```python
results = evaluator.evaluate_all(
    reference=["Reference translation."],
    source=["源文本。"],
    target_text=["My translation."],
)
# {'sacreBLEU': 45.2, 'chrF++': 62.1, 'COMET': 0.85, 'hypothesis_text': [...]}
```

### 纯语音评测

```python
evaluator = ModelEvaluator(use_comet=True, use_whisper=True)

results = evaluator.evaluate_all(
    reference=["Reference translation."],
    source=["源文本。"],
    target_speech="./my_audio/",
)
# {'sacreBLEU_ASR': 38.1, 'chrF++_ASR': 55.3, 'COMET_ASR': 0.78, 'hypothesis_ASR': [...]}
```

### 双模式评测（同时输入文本和语音）

```python
evaluator = ModelEvaluator(use_comet=True, use_whisper=True)

results = evaluator.evaluate_all(
    reference=["Reference translation."],
    source=["源文本。"],
    target_text=["My translation."],
    target_speech="./my_audio/",
)
# {
#     'sacreBLEU': 45.2,        'sacreBLEU_ASR': 38.1,
#     'chrF++': 62.1,           'chrF++_ASR': 55.3,
#     'COMET': 0.85,            'COMET_ASR': 0.78,
#     'hypothesis_text': [...],
#     'hypothesis_ASR': [...],
# }
```

> 📌 同时输入时，`target_text` 和 `target_speech` 必须是**同一批样本**的不同形式，数量需一致。

---

## 📁 使用内置数据集

```python
from multimetric_eval import ModelEvaluator, load_dataset

# 加载内置数据集（自动下载到 ./datasets/）
# 若有网络问题，可手动下载：
# https://github.com/sjtuayj/MultiMetric-Eval/releases/download/v0.1.0/zh-en-littleprince.zip
# 解压后将 zh-en-littleprince 文件夹保存至 ./datasets/
dataset = load_dataset("zh-en-littleprince")
```

### 方式1：传入文本列表

```python
evaluator = ModelEvaluator(use_comet=True)

results = evaluator.evaluate_dataset(
    dataset=dataset,
    target_text=["Translation 1", "Translation 2", ...],
)
```

### 方式2：传入 JSON / TXT 文件

```python
results = evaluator.evaluate_dataset(
    dataset=dataset,
    target_text="translations.json",  # 或 "translations.txt"
)
```

### 方式3：传入音频文件夹

```python
evaluator = ModelEvaluator(use_comet=True, use_whisper=True)

results = evaluator.evaluate_dataset(
    dataset=dataset,
    target_speech="./my_audio/",
)
```

### 方式4：同时传入文本和语音

```python
evaluator = ModelEvaluator(use_comet=True, use_whisper=True)

results = evaluator.evaluate_dataset(
    dataset=dataset,
    target_text=["Translation 1", "Translation 2", ...],
    target_speech="./my_audio/",
)
# 返回两组指标：sacreBLEU / sacreBLEU_ASR, chrF++ / chrF++_ASR, ...
```

---

## 📂 使用自定义数据集

```python
from multimetric_eval import ModelEvaluator

evaluator = ModelEvaluator(use_comet=True)

reference = ["Reference 1", "Reference 2"]
source = ["源文本1", "源文本2"]  # COMET 需要
```

### 纯文本评测

```python
# 传入列表
results = evaluator.evaluate(
    hypothesis=["Translation 1", "Translation 2"],
    reference=reference,
    source=source,
)

# 传入文件
results = evaluator.evaluate_file(
    hypothesis_file="translations.json",  # 或 .txt
    reference=reference,
    source=source,
)
```

### 纯语音评测

```python
evaluator = ModelEvaluator(use_comet=True, use_whisper=True)

results = evaluator.evaluate_audio_folder(
    audio_folder="./my_audio/",
    reference=reference,
    source=source,
)
```

### 双模式评测（统一接口）

```python
evaluator = ModelEvaluator(use_comet=True, use_whisper=True)

results = evaluator.evaluate_all(
    reference=reference,
    source=source,
    target_text=["Translation 1", "Translation 2"],  # 或文件路径
    target_speech="./my_audio/",
)
```

---

## 📄 输入文件格式

### JSON 文件（三种格式均支持）

**格式1：字典格式**
```json
{
    "hypothesis": [
        "Translation sentence 1.",
        "Translation sentence 2."
    ]
}
```

**格式2：对象数组格式**
```json
[
    {"id": "001", "hypothesis": "Translation sentence 1."},
    {"id": "002", "hypothesis": "Translation sentence 2."}
]
```

**格式3：纯字符串数组**
```json
[
    "Translation sentence 1.",
    "Translation sentence 2."
]
```

### TXT 文件

每行一句，空行自动忽略：

```text
Translation sentence 1.
Translation sentence 2.
```

### 音频文件夹

```
my_audio/
├── 001.wav
├── 002.wav
├── 003.mp3
└── 004.flac
```

- **支持格式**：`.wav`、`.mp3`、`.flac`
- **排序规则**：按文件名自动排序（确保与参考译文顺序一致）
- **命名建议**：使用数字前缀如 `001.wav`、`002.wav`

---

## ⚙️ 参数配置

### 评测器参数

```python
evaluator = ModelEvaluator(
    use_comet=True,                        # 启用 COMET（需要 source）
    use_bleurt=False,                      # 启用 BLEURT
    use_whisper=False,                     # 启用语音转文字
    comet_model="Unbabel/wmt22-comet-da",  # COMET 模型
    whisper_model="medium",                # tiny/base/small/medium/large
    bleurt_path=None,                      # BLEURT 本地模型路径（优先使用）
    bleurt_model=None,                     # BLEURT 在线模型名称（无本地路径时使用）
    device=None,                           # cuda/cuda:0/cuda:1/cpu，默认自动检测
)
```

| 参数 | 类型 | 默认值 | 说明 |
|------|------|--------|------|
| `use_comet` | bool | `True` | 启用 COMET 指标 |
| `use_bleurt` | bool | `False` | 启用 BLEURT 指标 |
| `use_whisper` | bool | `False` | 启用语音转文字 |
| `comet_model` | str | `"Unbabel/wmt22-comet-da"` | COMET 模型名称 |
| `whisper_model` | str | `"medium"` | Whisper 模型大小 |
| `bleurt_path` | str/None | `None` | BLEURT 本地模型路径（优先级最高） |
| `bleurt_model` | str/None | `None` | BLEURT 在线模型名称，默认 `lucadiliello/BLEURT-20` |
| `device` | str/None | `None` | 计算设备，支持 `cuda:N` 指定 GPU |

### BLEURT 模型加载优先级

```
bleurt_path（本地路径存在）  →  直接加载本地模型
        ↓ 不存在
bleurt_model（用户指定）     →  从 HuggingFace 下载
        ↓ 未指定
默认模型 lucadiliello/BLEURT-20  →  从 HuggingFace 下载
```

### 数据集参数

```python
dataset = load_dataset(
    name="zh-en-littleprince",   # 数据集名称
    cache_dir="./datasets",      # 缓存目录
    force_download=False,        # 强制重新下载
)
```

---

## 🖥️ GPU 指定

COMET、BLEURT、Whisper **统一运行在同一设备上**（CPU 或 GPU），通过 `device` 参数控制。

### 方式1：代码中指定

```python
# 使用第 0 号 GPU
evaluator = ModelEvaluator(device="cuda:0")

# 使用第 3 号 GPU
evaluator = ModelEvaluator(device="cuda:3")

# 强制使用 CPU
evaluator = ModelEvaluator(device="cpu")

# 自动选择（默认，有 GPU 用 GPU）
evaluator = ModelEvaluator()
```

### 方式2：命令行环境变量

```bash
# 只使用第 2 号 GPU
CUDA_VISIBLE_DEVICES=2 python my_eval.py

# 使用第 0 和第 1 号 GPU
CUDA_VISIBLE_DEVICES=0,1 python my_eval.py

# 禁用 GPU，强制 CPU
CUDA_VISIBLE_DEVICES="" python my_eval.py
```

### 方式3：脚本中使用 argparse

```python
import argparse
from multimetric_eval import ModelEvaluator

parser = argparse.ArgumentParser()
parser.add_argument("--device", type=str, default=None, help="指定GPU，如 cuda:0, cuda:1, cpu")
args = parser.parse_args()

evaluator = ModelEvaluator(device=args.device)
```

```bash
python my_eval.py --device cuda:2
```

---

## 📊 支持的指标

| 指标 | 说明 | 需要 source | 运行设备 | 需要额外安装 |
|------|------|-------------|----------|--------------|
| sacreBLEU | 标准 BLEU 分数 | ❌ | CPU | ❌（内置） |
| chrF++ | 字符级 F 分数 | ❌ | CPU | ❌（内置） |
| COMET | 神经网络评估 | ✅ | GPU/CPU | `pip install unbabel-comet` |
| BLEURT | Google BLEURT (PyTorch) | ❌ | GPU/CPU | `pip install git+https://github.com/lucadiliello/bleurt-pytorch.git` |

> 涉及语音输入（ASR）时，以上每个指标均会额外输出带 `_ASR` 后缀的版本。

---

## 📤 输出结果

### 纯文本输入

```python
{
    "sacreBLEU": 45.23,
    "chrF++": 62.15,
    "COMET": 0.8523,          # use_comet=True 时
    "BLEURT": 0.7234,         # use_bleurt=True 时
    "hypothesis_text": [...], # evaluate_all / evaluate_dataset 时返回
}
```

### 纯语音输入

```python
{
    "sacreBLEU_ASR": 38.12,
    "chrF++_ASR": 55.30,
    "COMET_ASR": 0.7823,
    "BLEURT_ASR": 0.6534,
    "hypothesis_ASR": [...],  # Whisper 转写结果
}
```

### 双模式输入

```python
{
    "sacreBLEU": 45.23,        "sacreBLEU_ASR": 38.12,
    "chrF++": 62.15,           "chrF++_ASR": 55.30,
    "COMET": 0.8523,           "COMET_ASR": 0.7823,
    "BLEURT": 0.7234,          "BLEURT_ASR": 0.6534,
    "hypothesis_text": [...],
    "hypothesis_ASR": [...],
}
```

---

## 📋 API 总结

| 方法 | 用途 | 输入方式 |
|------|------|----------|
| `evaluate()` | 纯文本评测 | `hypothesis` 列表 |
| `evaluate_file()` | 从文件评测 | JSON / TXT 文件路径 |
| `evaluate_audio_folder()` | 纯语音评测 | 音频文件夹路径 |
| `evaluate_all()` | 统一接口（自定义数据） | `target_text` 和/或 `target_speech` |
| `evaluate_dataset()` | 统一接口（内置数据集） | `target_text` 和/或 `target_speech` |

---

## 🔧 高级用法

### 使用上下文管理器（自动释放显存）

```python
with ModelEvaluator(use_comet=True) as evaluator:
    results = evaluator.evaluate(
        hypothesis=["Translation"],
        reference=["Reference"],
        source=["源文本"],
    )
# 退出 with 块后自动释放显存
```

### 手动释放显存

```python
evaluator = ModelEvaluator(use_comet=True, use_bleurt=True)
results = evaluator.evaluate(...)

# 评测完成后手动释放
evaluator.cleanup()
```

### 从本地 JSON 创建自定义数据集

```python
from multimetric_eval import create_dataset_from_json

# my_data.json 格式：
# [
#     {"id": "001", "source_text": "源文本1", "reference_text": "Ref 1"},
#     {"id": "002", "source_text": "源文本2", "reference_text": "Ref 2"}
# ]

dataset = create_dataset_from_json("./my_data.json")

results = evaluator.evaluate_dataset(
    dataset=dataset,
    target_text=["Translation 1", "Translation 2"],
)
```

### 查看可用数据集

```python
from multimetric_eval import list_datasets, get_dataset_info

print(list_datasets())
# ['zh-en-littleprince']

info = get_dataset_info("zh-en-littleprince")
print(info)
# {
#     'name': 'zh-en-littleprince',
#     'is_downloaded': True,
#     'num_samples': 54,
#     'audio_complete': True
# }
```

### 向后兼容（旧版参数）

```python
# 以下旧写法依然有效
results = evaluator.evaluate_dataset(
    dataset=dataset,
    hypothesis=["Translation 1", ...],   # 等同于 target_text
)

results = evaluator.evaluate_dataset(
    dataset=dataset,
    audio_folder="./my_audio/",          # 等同于 target_speech
)
```

---

## ❓ 常见问题

### Q: COMET 分数显示 -1.0？
A: 请确保传入了 `source` 参数，COMET 需要源文本。

### Q: CUDA out of memory？
A: 使用上下文管理器或手动调用 `evaluator.cleanup()` 释放显存。也可以通过 `device="cuda:N"` 指定空闲 GPU，或选择较小的模型（如 `BLEURT-20-D12`）。

### Q: 如何只使用基础指标？
A: 设置 `use_comet=False`，只计算 sacreBLEU 和 chrF++，无需下载任何模型。

### Q: 音频文件顺序不对？
A: 使用数字前缀命名，如 `001.wav`、`002.wav`，确保排序正确。

### Q: BLEURT 和 COMET 环境冲突？
A: **不会冲突**。本工具使用 PyTorch 版 BLEURT（[bleurt-pytorch](https://github.com/lucadiliello/bleurt-pytorch)），与 COMET 共享同一个 PyTorch 环境，所有模型统一运行在同一设备上。

### Q: 在中国大陆服务器模型下载失败？
A: 在代码最前面添加 HuggingFace 镜像：
```python
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
```

### Q: 如何指定使用哪块 GPU？
A: 代码中设置 `device="cuda:N"`，或命令行使用 `CUDA_VISIBLE_DEVICES=N python script.py`。

### Q: 离线环境如何使用 BLEURT？
A: 在联网机器上下载模型目录，传输到离线机器后使用 `bleurt_path` 参数指定本地路径。详见上方「BLEURT 安装」章节。

### Q: BLEURT-20 太大了怎么办？
A: 使用轻量版 `BLEURT-20-D12`：
```python
evaluator = ModelEvaluator(
    use_bleurt=True,
    bleurt_model="lucadiliello/BLEURT-20-D12",
)
```

---

## 🔄 版本更新

### v0.1.2（当前版本）

- **BLEURT 改用 PyTorch 版**：不再依赖 TensorFlow，彻底解决与 COMET 的环境冲突
- **BLEURT 支持 GPU**：与 COMET/Whisper 统一运行在同一设备上
- **BLEURT 支持在线下载**：新增 `bleurt_model` 参数，无需手动下载模型文件
- **BLEURT 分批推理**：自动分批处理（batch_size=32），避免显存溢出

### v0.1.1

- 新增三种评测模式（纯文本/纯语音/双模式）
- 新增 GPU 指定功能
- 新增向后兼容旧版参数

### v0.1.0

- 初始版本

---

## 📜 License

MIT License

---

## 🤝 Contributing

欢迎提交 Issue 和 Pull Request！

GitHub: [https://github.com/sjtuayj/MultiMetric-Eval](https://github.com/sjtuayj/MultiMetric-Eval)