refactor(models): accelerate evaluating using mp

ChenZiHong-Gavin · ChenZiHong-Gavin · commit 62f0f79a4887 · 2025-01-16T17:13:50.000+08:00
diff --git a/evaluate.py b/evaluate.py
@@ -3,9 +3,10 @@
 import os
 import json
 import argparse
-import torch
 import pandas as pd
 from dotenv import load_dotenv
+import torch
+import torch.multiprocessing as mp
 from models import LengthEvaluator, MTLDEvaluator, RewardEvaluator, TextPair, UniEvaluator
 from utils import logger, set_logger
 
@@ -51,15 +52,12 @@ def evaluate_uni(corpus, uni_model_name):
         model_name=uni_model_name
     )
     logger.info("Uni evaluator loaded with model %s", uni_model_name)
-    naturalness_scores = uni_evaluator.get_average_score(corpus, 'naturalness')
-    logger.info("Uni naturalness scores: %s", naturalness_scores)
-    coherence_scores = uni_evaluator.get_average_score(corpus, 'coherence')
-    logger.info("Uni coherence scores: %s", coherence_scores)
-    understandability_scores = uni_evaluator.get_average_score(corpus, 'understandability')
-    logger.info("Uni understandability scores: %s", understandability_scores)
+    uni_scores = uni_evaluator.get_average_score(corpus)
+    for key, value in uni_scores.items():
+        logger.info("Uni %s scores: %s", key, value)
     del uni_evaluator
     clean_gpu_cache()
-    return naturalness_scores, coherence_scores, understandability_scores
+    return uni_scores['naturalness'], uni_scores['coherence'], uni_scores['understandability']
 
 
 def clean_gpu_cache():
@@ -92,6 +90,8 @@ def clean_gpu_cache():
     results = []
 
     logger.info("Data loaded from %s", args.folder)
+    mp.set_start_method('spawn')
+
     for file in os.listdir(args.folder):
         if file.endswith('.json'):
             logger.info("Processing %s", file)
diff --git a/models/evaluate/reward_evaluator.py b/models/evaluate/reward_evaluator.py
@@ -1,37 +1,90 @@
 from dataclasses import dataclass
+from tqdm import tqdm
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
+import torch.multiprocessing as mp
 
-from models.evaluate.base_evaluator import BaseEvaluator
 from models.text.text_pair import TextPair
-from utils import create_event_loop
 
 
 @dataclass
-class RewardEvaluator(BaseEvaluator):
+class RewardEvaluator:
     """
     Reward Model Evaluator.
     OpenAssistant/reward-model-deberta-v3-large-v2: 分数范围为[-inf, inf]，越高越好
     """
     reward_name: str = "OpenAssistant/reward-model-deberta-v3-large-v2"
-    max_length: int = 1024
+    max_length: int = 2560
 
     def __post_init__(self):
-        self.rank_model = AutoModelForSequenceClassification.from_pretrained(self.reward_name)
-        self.tokenizer = AutoTokenizer.from_pretrained(self.reward_name)
+        self.num_gpus = torch.cuda.device_count()
 
-        self.rank_model.eval()
-        self.rank_model.to("cuda")
+    @staticmethod
+    def process_chunk(rank, pairs, reward_name, max_length, return_dict):
+        device = f'cuda:{rank}'
+        torch.cuda.set_device(rank)
 
-    async def evaluate_single(self, pair: TextPair) -> float:
-        loop = create_event_loop()
-        return await loop.run_in_executor(None, self._tokenize_and_rank, pair)
+        rank_model = AutoModelForSequenceClassification.from_pretrained(reward_name)
+        tokenizer = AutoTokenizer.from_pretrained(reward_name)
+        rank_model.to(device)
+        rank_model.eval()
 
-    def _tokenize_and_rank(self, pair: TextPair) -> float:
-        question, answer = pair.question, pair.answer
+        results = []
+        with torch.no_grad():
+            for pair in tqdm(pairs):
+                inputs = tokenizer(
+                    pair.question,
+                    pair.answer,
+                    return_tensors="pt",
+                    max_length=max_length,
+                    truncation=True
+                )
+                inputs = {k: v.to(device) for k, v in inputs.items()}
+                score = rank_model(**inputs).logits[0].item()
+                results.append(score)
 
-        # concatenate the question and answer
-        inputs = self.tokenizer(question, answer, return_tensors="pt", max_length=self.max_length, truncation=True)
-        inputs = {k: v.to("cuda") for k, v in inputs.items()}
+        return_dict[rank] = results
 
-        score = self.rank_model(**inputs).logits[0].item()
-        return score
+    def evaluate(self, pairs: list[TextPair]) -> list[float]:
+        chunk_size = len(pairs) // self.num_gpus
+        chunks = []
+        for i in range(self.num_gpus):
+            start = i * chunk_size
+            end = start + chunk_size
+            if i == self.num_gpus - 1:
+                end = len(pairs)
+            chunks.append(pairs[start:end])
+
+        # multi-process
+        manager = mp.Manager()
+        return_dict = manager.dict()
+        processes = []
+
+        for rank, chunk in enumerate(chunks):
+            p = mp.Process(
+                target=self.process_chunk,
+                args=(rank, chunk, self.reward_name, self.max_length, return_dict)
+            )
+            p.start()
+            processes.append(p)
+
+        for p in processes:
+            p.join()
+
+        # 合并结果
+        results = []
+        for rank in range(len(chunks)):
+            results.extend(return_dict[rank])
+
+        for p in processes:
+            if p.is_alive():
+                p.terminate()
+                p.join()
+
+        return results
+
+    def get_average_score(self, pairs: list[TextPair]) -> float:
+        """
+        Get the average score of a batch of texts.
+        """
+        return sum(self.evaluate(pairs)) / len(pairs)
diff --git a/models/evaluate/uni_evaluator.py b/models/evaluate/uni_evaluator.py
@@ -1,124 +1,146 @@
 # https://github.com/maszhongming/UniEval/tree/main
 
+from dataclasses import dataclass, field
+from tqdm import tqdm
 import torch
 from torch import nn
-from dataclasses import dataclass, field
-import asyncio
-from tqdm.asyncio import tqdm as tqdm_async
+import torch.multiprocessing as mp
 
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from models.evaluate.base_evaluator import BaseEvaluator
-from utils import create_event_loop
-from models.text.text_pair import TextPair
-
+from models import TextPair
+
+
+def _add_questions(dimension: str, question: str, answer: str):
+    if dimension == "naturalness":
+        cur_input = 'question: Is this a natural response in the dialogue? </s> response: ' + answer
+    elif dimension == "coherence":
+        cur_input = 'question: Is this a coherent response given the dialogue history? </s> response: ' \
+                    + answer + ' </s> dialogue history: ' + question
+    elif dimension == "understandability":
+        cur_input = 'question: Is this an understandable response in the dialogue? </s> response: ' + answer
+    else:
+        raise NotImplementedError(
+            'The input format for this dimension is still undefined. Please customize it first.')
+    return cur_input
 
 @dataclass
-class UniEvaluator(BaseEvaluator):
+class UniEvaluator:
     model_name: str = "MingZhong/unieval-sum"
     dimensions: list = field(default_factory=lambda: ['naturalness', 'coherence', 'understandability'])
-    max_length: int = 1024
+    max_length: int = 2560
 
     def __post_init__(self):
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-
-        self.model.eval()
-        self.model.to("cuda")
+        self.num_gpus = torch.cuda.device_count()
 
-        self.softmax = nn.Softmax(dim=1)
+    @staticmethod
+    def process_chunk(rank, pairs, model_name, max_length, dimension, return_dict):
+        device = f'cuda:{rank}'
+        torch.cuda.set_device(rank)
 
-        self.pos_id = self.tokenizer("Yes")["input_ids"][0]
-        self.neg_id = self.tokenizer("No")["input_ids"][0]
+        rank_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        rank_model.to(device)
+        rank_model.eval()
 
-    def evaluate(self, pairs: list[TextPair], dimension: str) -> list[float]:
-        """
-        Evaluate the text and return a score.
-        """
-        return create_event_loop().run_until_complete(self.async_evaluate(pairs, dimension))
+        softmax = nn.Softmax(dim=1)
 
-    async def async_evaluate(self, pairs: list[TextPair], dimension: str) -> list[float]:
-        semaphore = asyncio.Semaphore(self.max_concurrent)
-
-        async def evaluate_with_semaphore(pair):
-            async with semaphore:
-                return await self.evaluate_single(pair, dimension)
+        pos_id = tokenizer("Yes")["input_ids"][0]
+        neg_id = tokenizer("No")["input_ids"][0]
 
         results = []
-        for result in tqdm_async(
-                asyncio.as_completed([evaluate_with_semaphore(pair) for pair in pairs]),
-                total=len(pairs),
-        ):
-            results.append(await result)
-        return results
-
-    async def evaluate_single(self, pair: TextPair, dimension: str) -> float:
-        text = self._add_questions(dimension, pair.question, pair.answer)
-        loop = create_event_loop()
-        return await loop.run_in_executor(None, self._score, text)
-
-    def get_average_score(self, pairs: list[TextPair], dimension: str) -> float:
+        with torch.no_grad():
+            for pair in tqdm(pairs):
+                text = _add_questions(dimension, pair.question, pair.answer)
+
+                tgt = "No"
+
+                encoded_src = tokenizer(
+                    text,
+                    max_length=max_length,
+                    truncation=True,
+                    padding=True,
+                    return_tensors='pt'
+                )
+                encoded_tgt = tokenizer(
+                    tgt,
+                    max_length=max_length,
+                    truncation=True,
+                    padding=True,
+                    return_tensors='pt'
+                )
+
+                src_tokens = encoded_src['input_ids'].to(device)
+                src_mask = encoded_src['attention_mask'].to(device)
+
+                tgt_tokens = encoded_tgt['input_ids'].to(device)[:, 0].unsqueeze(-1)
+
+                output = rank_model(
+                    input_ids=src_tokens,
+                    attention_mask=src_mask,
+                    labels=tgt_tokens,
+                    use_cache = False
+                )
+
+                logits = output.logits.view(-1, rank_model.config.vocab_size)
+
+                pos_score = softmax(logits)[:, pos_id]  # Yes
+                neg_score = softmax(logits)[:, neg_id]
+                score = pos_score / (pos_score + neg_score)
+
+                results.append(score.item())
+
+        return_dict[rank] = results
+
+    def evaluate(self, pairs: list[TextPair]) -> list[dict]:
+        final_results = []
+        for dimension in self.dimensions:
+            chunk_size = len(pairs) // self.num_gpus
+            chunks = []
+            for i in range(self.num_gpus):
+                start = i * chunk_size
+                end = start + chunk_size
+                if i == self.num_gpus - 1:
+                    end = len(pairs)
+                chunks.append(pairs[start:end])
+
+            # multi-process
+            manager = mp.Manager()
+            return_dict = manager.dict()
+            processes = []
+
+            for rank, chunk in enumerate(chunks):
+                p = mp.Process(
+                    target=self.process_chunk,
+                    args=(rank, chunk, self.model_name, self.max_length, dimension, return_dict)
+                )
+                p.start()
+                processes.append(p)
+
+            for p in processes:
+                p.join()
+
+            # 合并结果
+            results = []
+            for rank in range(len(chunks)):
+                results.extend(return_dict[rank])
+
+            for p in processes:
+                if p.is_alive():
+                    p.terminate()
+                    p.join()
+
+            final_results.append({
+                dimension: results
+            })
+        return final_results
+
+    def get_average_score(self, pairs: list[TextPair]) -> dict:
         """
         Get the average score of a batch of texts.
         """
-        return sum(self.evaluate(pairs, dimension)) / len(pairs)
-
-    def _score(self, text: str) -> float:
-        """
-            Get scores for the given samples.
-            final_score = postive_score / (postive_score + negative_score)
-        """
-
-        # The implementation of "forward" in T5 still requires decoder_input_ids.
-        # Therefore, we construct a random one-word target sequence.
-        # The content of the target has no effect on the final scores.
-
-        tgt = "No"
-
-        with torch.no_grad():
-            encoded_src = self.tokenizer(
-                text,
-                max_length=self.max_length,
-                truncation=True,
-                padding=True,
-                return_tensors='pt'
-            )
-            encoded_tgt = self.tokenizer(
-                tgt,
-                max_length=self.max_length,
-                truncation=True,
-                padding=True,
-                return_tensors='pt'
-            )
-
-            src_tokens = encoded_src['input_ids'].to("cuda")
-            src_mask = encoded_src['attention_mask'].to("cuda")
-
-            tgt_tokens = encoded_tgt['input_ids'].to("cuda")[:, 0].unsqueeze(-1)
-
-            output = self.model(
-                input_ids=src_tokens,
-                attention_mask=src_mask,
-                labels=tgt_tokens
-            )
-
-            logits = output.logits.view(-1, self.model.config.vocab_size)
-
-            pos_score = self.softmax(logits)[:, self.pos_id]  # Yes
-            neg_score = self.softmax(logits)[:, self.neg_id]
-
-            score = pos_score / (pos_score + neg_score)
-
-        return score.item()
-
-    def _add_questions(self, dimension: str, question: str, answer: str):
-        if dimension == "naturalness":
-            cur_input = 'question: Is this a natural response in the dialogue? </s> response: ' + answer
-        elif dimension == "coherence":
-            cur_input = 'question: Is this a coherent response given the dialogue history? </s> response: ' \
-                        + answer + ' </s> dialogue history: ' + question
-        elif dimension == "understandability":
-            cur_input = 'question: Is this an understandable response in the dialogue? </s> response: ' + answer
-        else:
-            raise NotImplementedError(
-                'The input format for this dimension is still undefined. Please customize it first.')
-        return cur_input
+        results = self.evaluate(pairs)
+        final_results = {}
+        for result in results:
+            for key, value in result.items():
+                final_results[key] = sum(value) / len(value)
+        return final_results
