fix: eval

Former-commit-id: 0a964e3f26e607fa7a85bbed9bff7ff13c13299a

evaluation/evaluation.py

@@ -1,57 +1,56 @@
 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
+from alignment import alignment
+from scores.multi_scores import multi_scores
 
 class Evaluator:
-    def __init__(self, src_path, pred_path, gt_path, eval_path, conclusion_path):
-        self.src_path = src_path
+    def __init__(self, pred_path, gt_path, eval_path, res_path):
         self.pred_path = pred_path
         self.gt_path = gt_path
         self.eval_path = eval_path
-        self.conclusion_path = conclusion_path
+        self.res_path = res_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
+        for (pred_s, gt_s) in aligned_srt:
+            print("pred_s.source_text: ", pred_s.source_text)
+            print("pred_s.translation: ", pred_s.translation)
+            print("gt_s.source_text: ", gt_s.translation)
+
+            scores_dict = scorer.get_scores(pred_s.source_text, pred_s.translation, gt_s.translation)
+            scores_dict['Source'] = pred_s.source_text
+            scores_dict['Prediction'] = pred_s.translation
+            scores_dict['Ground Truth'] = gt_s.translation
             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'])
+        eval_df.to_csv(self.eval_path, index=False, columns=['Source', 'Prediction', 'Ground Truth', 'bleu_score', 'comet_score', 'llm_score', 'llm_explanation'])
 
         # Get average scores
-        avg_llm = eval_df['llm'].mean()
-        avg_bleu = eval_df['bleu'].mean()
-        avg_comet = eval_df['comet'].mean()
+        avg_llm = eval_df['llm_score'].mean()
+        avg_bleu = eval_df['bleu_score'].mean()
+        avg_comet = eval_df['comet_score'].mean()
 
-        conclusion_data = {
+        res_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)
+        res_df = pd.DataFrame(res_data)
+        res_df.to_csv(self.res_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')
+    parser.add_argument('-bi_path', default='evaluation/test5_tiny/test5_bi.srt', help='Path to predicted SRT file')
+    parser.add_argument('-zh_path', default='evaluation/test5_tiny/test5_gt.srt', help='Path to ground truth SRT file')
+    parser.add_argument('-eval_output', default='evaluation/test5_tiny/eval.csv', help='Path to eval CSV file')
+    parser.add_argument('-res_output', default='evaluation/test5_tiny/res.csv', help='Path to result CSV file')
     args = parser.parse_args()
 
-    evaluator = Evaluator(args.src, args.pred, args.gt, args.eval, args.conclusion)
+    evaluator = Evaluator(args.bi_path, args.zh_path, args.eval_output, args.res_output)
     evaluator.eval()