darpanaswal commited on
Commit
48c5057
·
verified ·
1 Parent(s): 61330e7

Update metrics.py

Browse files
Files changed (1) hide show
  1. metrics.py +9 -5
metrics.py CHANGED
@@ -5,18 +5,22 @@ from bert_score import score as bert_score
5
 
6
  def compute_scores(predictions, references):
7
  scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True)
8
- scores = {"ROUGE-1": [], "ROUGE-2": [], "ROUGE-L": []}
9
-
 
10
  for pred, ref in zip(predictions, references):
11
  rouge_scores = scorer.score(pred, ref)
12
  scores["ROUGE-1"].append(rouge_scores["rouge1"].fmeasure)
13
  scores["ROUGE-2"].append(rouge_scores["rouge2"].fmeasure)
14
  scores["ROUGE-L"].append(rouge_scores["rougeL"].fmeasure)
15
-
 
 
 
 
16
  return {key: np.mean(value) for key, value in scores.items()}
17
 
18
  def save_scores(scores, model_name, experiment_type, dataset_name):
19
  with open("rouge_results.csv", mode="a", newline="") as file:
20
  writer = csv.writer(file)
21
- writer.writerow([model_name, experiment_type, dataset_name, scores["ROUGE-1"], scores["ROUGE-2"], scores["ROUGE-L"]])
22
-
 
5
 
6
  def compute_scores(predictions, references):
7
  scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True)
8
+ scores = {"ROUGE-1": [], "ROUGE-2": [], "ROUGE-L": [], "BERT-F1": []}
9
+
10
+ # Compute ROUGE scores
11
  for pred, ref in zip(predictions, references):
12
  rouge_scores = scorer.score(pred, ref)
13
  scores["ROUGE-1"].append(rouge_scores["rouge1"].fmeasure)
14
  scores["ROUGE-2"].append(rouge_scores["rouge2"].fmeasure)
15
  scores["ROUGE-L"].append(rouge_scores["rougeL"].fmeasure)
16
+
17
+ # Compute BERTScore F1
18
+ P, R, F1 = bert_score(predictions, references, lang="en", rescale_with_baseline=True)
19
+ scores["BERT-F1"].extend(F1.tolist())
20
+
21
  return {key: np.mean(value) for key, value in scores.items()}
22
 
23
  def save_scores(scores, model_name, experiment_type, dataset_name):
24
  with open("rouge_results.csv", mode="a", newline="") as file:
25
  writer = csv.writer(file)
26
+ writer.writerow([model_name, experiment_type, dataset_name, scores["ROUGE-1"], scores["ROUGE-2"], scores["ROUGE-L"], scores["BERT-F1"]])