Merge branch 'evaluation' of github.com:project-kxkg/project-t into evaluation

llm and other scores fix


Former-commit-id: 878c36401ab6a5cfef8be2bd24f58dab749c08e2

evaluation/alignment.py

@@ -3,69 +3,87 @@ import numpy as np
 sys.path.append('../src')
 from srt_util.srt import SrtScript
 
+
+# Helper method
+# Align sub anchor segment pair via greedy approach
+# Input: anchor segment, SRT segments, output array of sub, index of current sub
+# Output: updated index of sub
 def procedure(anchor, subsec, S_arr, subidx):
     cache_idx = 0
-    while subidx != cache_idx:
+    while subidx != cache_idx:  # Terminate when alignment stablizes
         cache_idx = subidx
-        if subidx >= len(subsec):
+        # if sub segment runs out during the loop, terminate
+        if subidx >= len(subsec): 
             break
         sub = subsec[subidx]
         if anchor.end < sub.start:
             continue
+        # If next sub has a heavier overlap compartment, add to current alignment
         if (anchor.start <= sub.start) and (sub.end <= anchor.end) or anchor.end - sub.start > sub.end - anchor.start:
-            S_arr[-1] += sub.source_text
+            S_arr[-1] += sub#.source_text
             subidx += 1
-    return subidx - 1
 
+    return subidx - 1  # Reset last invalid update from loop
+
+
+# Input: path1, path2
+# Output: aligned array of SRTsegment corresponding to path1 path2
+# Note: Modify comment with .source_text to get output array with string only
 def alignment(pred_path, gt_path):
     pred = SrtScript.parse_from_srt_file(pred_path).segments
     gt = SrtScript.parse_from_srt_file(gt_path).segments
     pred_arr, gt_arr = [], []
-    idx_p, idx_t = 0, 0
+    idx_p, idx_t = 0, 0  # idx_p: current index of pred segment, idx_t for ground truth
 
     while idx_p < len(pred) or idx_t < len(gt):
+        # Check if one srt file runs out while reading
         ps = pred[idx_p] if idx_p < len(pred) else None
         gs = gt[idx_t] if idx_t < len(gt) else None
-
+        
         if not ps:
-            gt_arr.append(gs.source_text)
+            # If ps runs out, align gs segment with filler one by one
+            gt_arr.append(gs)#.source_text
             pred_arr.append('')
             idx_t += 1
             continue
 
         if not gs:
-            pred_arr.append(ps.source_text)
+            # If gs runs out, align ps segment with filler one by one
+            pred_arr.append(ps)#.source_text 
             gt_arr.append('')
             idx_p += 1
             continue
 
         ps_dur = ps.end - ps.start
         gs_dur = gs.end - gs.start
-
+        
+        # Check for duration to decide anchor and sub
         if ps_dur <= gs_dur:
+            # Detect segment with no overlap
             if ps.end < gs.start:
-                pred_arr.append(ps.source_text)
-                gt_arr.append('')
-                idx_t -= 1
+                pred_arr.append(ps)#.source_text
+                gt_arr.append('')  # append filler
+                idx_t -= 1  # reset ground truth index
             else:
-                gt_arr.append(gs.source_text)
+                gt_arr.append(gs)#.source_text
                 if gs.end >= ps.start:
-                    pred_arr.append(ps.source_text)
+                    pred_arr.append(ps)#.source_text
                     idx_p = procedure(gs, pred, pred_arr, idx_p + 1)
-                else:
+                else:  # filler pairing
                     pred_arr.append('')
                     idx_p -= 1
         else:
+            # same overlap checking procedure
             if gs.end < ps.start:
-                gt_arr.append(gs.source_text)
-                pred_arr.append('')
-                idx_p -= 1
+                gt_arr.append(gs)#.source_text
+                pred_arr.append('')  # filler
+                idx_p -= 1  # reset
             else:
-                pred_arr.append(ps.source_text)
+                pred_arr.append(ps)#.source_text
                 if ps.end >= gs.start:
-                    gt_arr.append(gs.source_text)
+                    gt_arr.append(gs)#.source_text
                     idx_t = procedure(ps, gt, gt_arr, idx_t + 1)
-                else:
+                else:  # filler pairing
                     gt_arr.append('')
                     idx_t -= 1
 

evaluation/evaluation.py

@@ -0,0 +1,57 @@
+import argparse
+import pandas as pd
+from evaluation.alignment import alignment
+from evaluation.scores.multi_scores import multi_scores
+from src.srt_util.srt import SrtScript
+
+class Evaluator:
+    def __init__(self, src_path, pred_path, gt_path, eval_path, conclusion_path):
+        self.src_path = src_path
+        self.pred_path = pred_path
+        self.gt_path = gt_path
+        self.eval_path = eval_path
+        self.conclusion_path = conclusion_path
+
+    def eval(self):
+        # Align two SRT files
+        aligned_srt = alignment(self.pred_path, self.gt_path)
+
+        # Parse src
+        src_s = [s.source_text for s in SrtScript.parse_from_srt_file(self.src_path).segments]
+
+        # Get sentence scores
+        scorer = multi_scores()
+        result_data = []
+        for ((prd_s, gt_s), src_s) in zip(aligned_srt, src_s):
+            scores_dict = scorer.get(src_s, prd_s, gt_s)
+            scores_dict['Prediction'] = prd_s
+            scores_dict['Ground Truth'] = gt_s
+            result_data.append(scores_dict)
+
+        eval_df = pd.DataFrame(result_data)
+        eval_df.to_csv(self.output_path, index=False, columns=['Prediction', 'Ground Truth', 'llm', 'bleu', 'comet'])
+
+        # Get average scores
+        avg_llm = eval_df['llm'].mean()
+        avg_bleu = eval_df['bleu'].mean()
+        avg_comet = eval_df['comet'].mean()
+
+        conclusion_data = {
+            'Metric': ['Avg LLM', 'Avg BLEU', 'Avg COMET'],
+            'Score': [avg_llm, avg_bleu, avg_comet]
+        }
+        conclusion_df = pd.DataFrame(conclusion_data)
+        conclusion_df.to_csv(self.conclusion_path, index=False)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Evaluate SRT files.')
+    parser.add_argument('-src', default='test/short_src', help='Path to source SRT file')
+    parser.add_argument('-pred', default='test/short_pred', help='Path to predicted SRT file')
+    parser.add_argument('-gt', default='test/short_gt', help='Path to ground truth SRT file')
+    parser.add_argument('-eval', default='eval.csv', help='Path to output CSV file')
+    parser.add_argument('-conclusion', default='conclusion.csv', help='Path to conclusion CSV file')
+    args = parser.parse_args()
+
+    evaluator = Evaluator(args.src, args.pred, args.gt, args.eval, args.conclusion)
+    evaluator.eval()
+

src/srt_util/srt.py

@@ -50,7 +50,10 @@ class SrtSegment(object):
             self.start = int(start_list[0]) * 3600 + int(start_list[1]) * 60 + int(start_list[2]) + self.start_ms / 100
             end_list = self.end_time_str.split(',')[0].split(':')
             self.end = int(end_list[0]) * 3600 + int(end_list[1]) * 60 + int(end_list[2]) + self.end_ms / 100
-            self.translation = ""
+            if len(args[0]) < 5:
+                self.translation = ""
+            else:
+                self.translation = args[0][3]
 
     def merge_seg(self, seg):
         """
@@ -105,10 +108,16 @@ class SrtScript(object):
     def parse_from_srt_file(cls, path: str):
         with open(path, 'r', encoding="utf-8") as f:
             script_lines = [line.rstrip() for line in f.readlines()]
-
+        bilingual = False
+        if script_lines[2] != '' and script_lines[3] != '':
+            bilingual = True
         segments = []
-        for i in range(0, len(script_lines), 4):
-            segments.append(list(script_lines[i:i + 4]))
+        if bilingual:
+            for i in range(0, len(script_lines), 5):
+                segments.append(list(script_lines[i:i + 5]))
+        else:
+            for i in range(0, len(script_lines), 4):
+                segments.append(list(script_lines[i:i + 4]))
 
         return cls(segments)