add Evaluator class

Former-commit-id: 6ba881680089ff333c7c47bf3a369f8655b1a511

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()
+