refactor: bleu score function name (#17)

shenxiangzhuang · Apr 23, 2024 · 029da9b · 029da9b
1 parent 4594c3d
commit 029da9b
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 19 deletions.
diff --git a/python/bleuscore/__init__.py b/python/bleuscore/__init__.py
@@ -4,5 +4,5 @@
 __all__ = [
     "tokenizer_regex",
     "tokenizer_13a",
-    "compute_bleu",
+    "compute",
 ]
diff --git a/src/bleu.rs b/src/bleu.rs
@@ -15,7 +15,7 @@ pub struct BleuScore {
 }
 
 
-pub fn bleu_score(
+pub fn compute_score(
     references: Vec<Vec<String>>,
     predictions: Vec<String>,
     max_order: usize,
@@ -115,14 +115,14 @@ pub fn bleu_score(
 
 #[cfg(test)]
 mod test {
-    use crate::bleu::{bleu_score};
+    use crate::bleu::{compute_score};
     #[test]
     fn test_bleu() {
         let reference_corpus: Vec<Vec<String>> = vec![vec!["Hello, World!".to_string()]];
         let translation_corpus: Vec<String> = vec!["Yellow, World!".to_string()];
         let max_order: usize = 4;
         let smooth: bool = true;
-        let res = bleu_score(reference_corpus, translation_corpus, max_order, smooth);
+        let res = compute_score(reference_corpus, translation_corpus, max_order, smooth);
         // (0.668740304976422, [0.8, 0.75, 0.6666666666666666, 0.5], 1.0, 1.0, 4, 4)
         println!("BLEU: {:?}", res);
         assert!((res.bleu - 0.668740304976422).abs() < 1e-10);

diff --git a/src/lib.rs b/src/lib.rs
@@ -2,7 +2,7 @@ mod tokenizer;
 pub use crate::tokenizer::{Tokenizer, Tokenizer13a, TokenizerRegex};
 mod ngram;
 mod bleu;
-pub use crate::bleu::{BleuScore, bleu_score};
+pub use crate::bleu::{BleuScore, compute_score};
 
 use pyo3::prelude::*;
 use pyo3::types::IntoPyDict;
@@ -23,13 +23,13 @@ fn tokenizer_13a(line: &str) -> PyResult<Vec<String>> {
 }
 
 #[pyfunction]
-fn compute_bleu(
+fn compute(
     references: Vec<Vec<String>>,
     predictions: Vec<String>,
     max_order: usize,
     smooth: bool,
 ) -> PyResult<PyObject> {
-    let bleu = bleu::bleu_score(references, predictions, max_order, smooth);
+    let bleu = compute_score(references, predictions, max_order, smooth);
     Python::with_gil(|py| {
         let bleu_dict = [
             ("bleu", bleu.bleu.to_object(py)), 
@@ -50,6 +50,6 @@ fn compute_bleu(
 fn bleuscore(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(tokenizer_regex, m)?)?;
     m.add_function(wrap_pyfunction!(tokenizer_13a, m)?)?;
-    m.add_function(wrap_pyfunction!(compute_bleu, m)?)?;
+    m.add_function(wrap_pyfunction!(compute, m)?)?;
     Ok(())
 }
diff --git a/tests/test_bleu_score.py b/tests/test_bleu_score.py
@@ -1,5 +1,4 @@
 import datetime
-from typing import List
 
 from hypothesis import given, settings
 from hypothesis import strategies as st
@@ -23,11 +22,11 @@ def local_py_bleu_result_compare(references, predictions, max_order, smooth):
 
     # same order with py_result, but here return a dict rather than list
     # bleu, precisions, bp, ratio, translation_length, reference_length
-    rust_result = bleuscore.compute_bleu(references=references,
-                                         predictions=predictions,
-                                         max_order=max_order,
-                                         smooth=smooth,
-                                         )
+    rust_result = bleuscore.compute(references=references,
+                                    predictions=predictions,
+                                    max_order=max_order,
+                                    smooth=smooth,
+                                    )
 
     for i, py_precision in enumerate(py_precisions):
         assert number_close(py_precision, rust_result["precisions"][i])
@@ -89,11 +88,11 @@ def hf_bleu_result_compare(references, predictions, max_order, smooth, debug=Fal
                                 max_order=max_order,
                                 smooth=smooth)
 
-    bleuscore_result = bleuscore.compute_bleu(references=references,
-                                              predictions=predictions,
-                                              max_order=max_order,
-                                              smooth=smooth,
-                                              )
+    bleuscore_result = bleuscore.compute(references=references,
+                                         predictions=predictions,
+                                         max_order=max_order,
+                                         smooth=smooth,
+                                         )
     if debug:
         print(f"hf_result: {hf_result}\n"
               f"bleuscore_result: {bleuscore_result}")