# 📊 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++ + COMET）
pip install multimetric-eval

# 可选依赖
pip install unbabel-comet    # COMET 指标
pip install openai-whisper   # 语音转文字
pip install bleurt           # BLEURT 指标
# 需要额外安装内容请见以下“支持的指标”部分表格
```

---

## 📖 快速开始

```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}
```

---

## 📁 使用内置数据集

```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")

# 初始化评测器
evaluator = ModelEvaluator(use_comet=True)
```

### 方式1：传入列表

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

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

```python
results = evaluator.evaluate_dataset(
    dataset=dataset,
    hypothesis="translations.json",
)
```

### 方式3：传入 TXT 文件

```python
results = evaluator.evaluate_dataset(
    dataset=dataset,
    hypothesis="translations.txt",
)
```

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

```python
# 需要启用 Whisper
evaluator = ModelEvaluator(use_comet=True, use_whisper=True)

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

---

## 📂 使用自定义数据集

```python
from multimetric_eval import ModelEvaluator

evaluator = ModelEvaluator(use_comet=True)

# 准备参考数据
reference = ["Reference 1", "Reference 2"]
source = ["源文本1", "源文本2"]  # COMET 需要
```

### 方式1：传入列表

```python
results = evaluator.evaluate(
    hypothesis=["Translation 1", "Translation 2"],
    reference=reference,
    source=source,
)
```

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

```python
results = evaluator.evaluate_file(
    hypothesis_file="translations.json",
    reference=reference,
    source=source,
)
```

### 方式3：传入 TXT 文件

```python
results = evaluator.evaluate_file(
    hypothesis_file="translations.txt",
    reference=reference,
    source=source,
)
```

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

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

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

---

## 📄 输入文件格式

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

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

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

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

### TXT 文件

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

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

### 音频文件夹

```
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 模型路径
    device=None,                           # cuda/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` | BLEURT 模型本地路径 |
| `device` | str | `None` | 计算设备，自动检测 GPU |

### 数据集参数

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

---

## 🎯 常用场景

### 场景1：快速评测（只用 BLEU + chrF++）

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

results = evaluator.evaluate(
    hypothesis=["My translation"],
    reference=["Reference translation"],
)
# {'sacreBLEU': 45.23, 'chrF++': 62.15}
```

### 场景2：完整评测（全部指标）

```python
evaluator = ModelEvaluator(
    use_comet=True,
    use_bleurt=True,
    bleurt_path="./model/BLEURT-20",
)

results = evaluator.evaluate(
    hypothesis=["My translation"],
    reference=["Reference translation"],
    source=["源文本"],
)
# {'sacreBLEU': 45.23, 'chrF++': 62.15, 'COMET': 0.85, 'BLEURT': 0.72}
```

### 场景3：语音评测

```python
evaluator = ModelEvaluator(
    use_comet=True,
    use_whisper=True,
    whisper_model="large",  # 更高精度
)

results = evaluator.evaluate_audio_folder(
    audio_folder="./speech_outputs/",
    reference=["Reference 1", "Reference 2"],
    source=["源文本1", "源文本2"],
)

# 查看 ASR 转写结果
print(results["hypothesis"])
```

### 场景4：强制使用 CPU

```python
evaluator = ModelEvaluator(
    use_comet=True,
    device="cpu",
)
```

---

## 📊 支持的指标

| 指标 | 说明 | 需要 source | 需要额外安装 |
|------|------|-------------|--------------|
| sacreBLEU | 标准 BLEU 分数 | ❌ | ❌ |
| chrF++ | 字符级 F 分数 | ❌ | ❌ |
| COMET | 神经网络评估 | ✅ | `unbabel-comet` |
| BLEURT | Google BLEURT | ❌ | `bleurt` + 模型文件 |

---

## 📤 输出结果

```python
results = evaluator.evaluate(...)

print(results)
# {
#     "sacreBLEU": 45.23,      # 始终返回
#     "chrF++": 62.15,         # 始终返回
#     "COMET": 0.8523,         # use_comet=True 时返回
#     "BLEURT": 0.7234,        # use_bleurt=True 时返回
#     "hypothesis": [...],     # 音频输入时返回转写结果
# }
```

---

## 📋 输入格式支持总结

| 格式 | 方法 | 示例 |
|------|------|------|
| Python 列表 | `evaluate()` | `["sent1", "sent2"]` |
| JSON 文件 | `evaluate_file()` | `"translations.json"` |
| TXT 文件 | `evaluate_file()` | `"translations.txt"` |
| 音频文件夹 | `evaluate_audio_folder()` | `"./audio/"` |

---

## 🔧 高级用法

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

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

### 从本地 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,
    hypothesis=["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': 100,
#     'audio_complete': True
# }
```

---

## ❓ 常见问题

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

### Q: CUDA out of memory？
A: 使用上下文管理器或手动调用 `evaluator.cleanup()` 释放显存。

### Q: 如何只使用基础指标？
A: 设置 `use_comet=False`，只计算 BLEU 和 chrF++。

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

---

## 📜 License

MIT License

---

## 🤝 Contributing

欢迎提交 Issue 和 Pull Request！