flashrag-dev 0.1.4.dev20250208__tar.gz → 0.1.4.dev20250209__tar.gz

{flashrag_dev-0.1.4.dev20250208 → flashrag_dev-0.1.4.dev20250209}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag_dev
-Version: 0.1.4.dev20250208
+Version: 0.1.4.dev20250209
 Summary: A library for efficient Retrieval-Augmented Generation research
 Home-page: https://github.com/RUC-NLPIR/FlashRAG
 Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou

{flashrag_dev-0.1.4.dev20250208 → flashrag_dev-0.1.4.dev20250209}/flashrag/evaluator/evaluator.py

@@ -15,7 +15,7 @@ class Evaluator:
 
         self.avaliable_metrics = self._collect_metrics()
 
-        self.metric_class = {}
+        self.metric_class = {}  
         for metric in self.metrics:
             if metric in self.avaliable_metrics:
                 self.metric_class[metric] = self.avaliable_metrics[metric](self.config)

{flashrag_dev-0.1.4.dev20250208 → flashrag_dev-0.1.4.dev20250209}/flashrag/evaluator/metrics.py

@@ -307,6 +307,8 @@ class Rouge_Score(BaseMetric):
         return output
 
 
+
+
 class Rouge_1(Rouge_Score):
     metric_name = "rouge-1"
 
@@ -364,6 +366,99 @@ class Rouge_L(Rouge_Score):
         return {"rouge-l": score}, metric_score_list
 
 
+
+class ZH_Rouge_Score(BaseMetric):
+    metric_name = "zh_rouge_score"
+    cached_scores = {}
+    
+    def __init__(self, config):
+        super().__init__(config)
+        from rouge_chinese import Rouge
+
+        self.scorer = Rouge()
+
+    def calculate_rouge(self, pred, golden_answers):
+        import jieba
+        if (pred, tuple(golden_answers)) in self.cached_scores:
+            return self.cached_scores[(pred, tuple(golden_answers))]
+        output = {}
+        pred = ' '.join(jieba.cut(pred))
+        for answer in golden_answers:
+            answer = ' '.join(jieba.cut(answer))
+            scores = self.scorer.get_scores(pred, answer)
+            for key in ["rouge-1", "rouge-2", "rouge-l"]:
+                if key not in output:
+                    output[key] = []
+                output[key].append(scores[0][key]["f"])
+        for k, v in output.items():
+            output[k] = max(v)
+
+        self.cached_scores[(pred, tuple(golden_answers))] = output
+        return output
+
+
+
+
+class ZH_Rouge_1(ZH_Rouge_Score):
+    metric_name = "zh_rouge-1"
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+
+    def calculate_metric(self, data):
+        golden_answers_list = self.get_dataset_answer(data)
+        pred_list = data.pred
+
+        metric_score_list = [
+            self.calculate_rouge(pred, golden_answers)["rouge-1"]
+            for pred, golden_answers in zip(pred_list, golden_answers_list)
+        ]
+        score = sum(metric_score_list) / len(metric_score_list)
+
+        return {"zh_rouge-1": score}, metric_score_list
+
+
+class ZH_Rouge_2(ZH_Rouge_Score):
+    metric_name = "zh_rouge-2"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+    def calculate_metric(self, data):
+        golden_answers_list = self.get_dataset_answer(data)
+        pred_list = data.pred
+
+        metric_score_list = [
+            self.calculate_rouge(pred, golden_answers)["rouge-2"]
+            for pred, golden_answers in zip(pred_list, golden_answers_list)
+        ]
+        score = sum(metric_score_list) / len(metric_score_list)
+
+        return {"zh_rouge-2": score}, metric_score_list
+
+
+class ZH_Rouge_L(ZH_Rouge_Score):
+    metric_name = "zh_rouge-l"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+    def calculate_metric(self, data):
+        golden_answers_list = self.get_dataset_answer(data)
+        pred_list = data.pred
+
+        metric_score_list = [
+            self.calculate_rouge(pred, golden_answers)["rouge-l"]
+            for pred, golden_answers in zip(pred_list, golden_answers_list)
+        ]
+        score = sum(metric_score_list) / len(metric_score_list)
+
+        return {"zh_rouge-l": score}, metric_score_list
+
+
+
+
 class BLEU(BaseMetric):
     metric_name = "bleu"
 

{flashrag_dev-0.1.4.dev20250208 → flashrag_dev-0.1.4.dev20250209}/flashrag/generator/openai_generator.py

@@ -4,11 +4,20 @@ from copy import deepcopy
 import warnings
 from tqdm import tqdm
 import numpy as np
-
+import threading
 import asyncio
 from openai import AsyncOpenAI, AsyncAzureOpenAI
 import tiktoken
 
+_background_loop = None
+
+def get_background_loop():
+    global _background_loop
+    if _background_loop is None:
+        _background_loop = asyncio.new_event_loop()
+        t = threading.Thread(target=lambda: _background_loop.run_forever(), daemon=True)
+        t.start()
+    return _background_loop
 
 class OpenaiGenerator:
     """Class for api-based openai models"""
@@ -54,32 +63,35 @@ class OpenaiGenerator:
     def update_additional_setting(self):
         pass
     
-    async def get_response(self, input: List, **params):
-        response = await self.client.chat.completions.create(model=self.model_name, messages=input, **params)
+    async def _get_response(self, messages: List, **params):
+        response = await self.client.chat.completions.create(
+            model=self.model_name, messages=messages, **params
+        )
+        if not response.choices:
+            raise ValueError("No choices returned from API.")
         return response.choices[0]
 
-    async def get_batch_response(self, input_list: List[List], batch_size, **params):
-        total_input = [self.get_response(input, **params) for input in input_list]
-        all_result = []
-        for idx in tqdm(range(0, len(input_list), batch_size), desc="Generation process: "):
-            batch_input = total_input[idx : idx + batch_size]
-            batch_result = await asyncio.gather(*batch_input)
-            all_result.extend(batch_result)
-
-        return all_result
-
-    def generate(self, input_list: List[List], batch_size=None, return_scores=False, **params) -> List[str]:
-        # deal with single input
-        if len(input_list) == 1 and isinstance(input_list[0], dict):
+    async def _get_batch_response(self, input_list: List[List], batch_size, **params):
+        tasks = [self._get_response(messages, **params) for messages in input_list]
+        all_results = []
+        for idx in tqdm(range(0, len(tasks), batch_size), desc="Generation process: "):
+            batch_tasks = tasks[idx: idx + batch_size]
+            batch_results = await asyncio.gather(*batch_tasks)
+            all_results.extend(batch_results)
+        return all_results
+
+    async def _generate_async(self, input_list: List, batch_size=None, return_scores=False, **params) -> List[str]:
+        if isinstance(input_list, dict):
+            input_list = [[input_list]]
+        elif isinstance(input_list[0], dict):
             input_list = [input_list]
+
         if batch_size is None:
             batch_size = self.batch_size
 
-        # deal with generation params
         generation_params = deepcopy(self.generation_params)
         generation_params.update(params)
-        if "do_sample" in generation_params:
-            generation_params.pop("do_sample")
+        generation_params.pop("do_sample", None)
 
         max_tokens = params.pop("max_tokens", None) or params.pop("max_new_tokens", None)
         if max_tokens is not None:
@@ -91,30 +103,26 @@ class OpenaiGenerator:
         generation_params.pop("max_new_tokens", None)
 
         if return_scores:
-            if generation_params.get("logprobs") is not None:
-                generation_params["logprobs"] = True
-                warnings.warn("Set logprobs to True to get generation scores.")
-            else:
-                generation_params["logprobs"] = True
-
-        if generation_params.get("n") is not None:
-            generation_params["n"] = 1
-            warnings.warn("Set n to 1. It can minimize costs.")
-        else:
-            generation_params["n"] = 1
+            generation_params["logprobs"] = True
+            warnings.warn("Set logprobs to True to get generation scores.")
+
 
-        loop = asyncio.get_event_loop()
-        result = loop.run_until_complete(self.get_batch_response(input_list, batch_size, **generation_params))
+        results = await self._get_batch_response(input_list, batch_size, **generation_params)
 
-        # parse result into response text and logprob
+        response_texts = []
         scores = []
-        response_text = []
-        for res in result:
-            response_text.append(res.message.content)
+        for res in results:
+            response_texts.append(res.message.content)
             if return_scores:
-                score = np.exp(list(map(lambda x: x.logprob, res.logprobs.content)))
+                score = np.exp([item.logprob for item in res.logprobs.content])
                 scores.append(score)
-        if return_scores:
-            return response_text, scores
-        else:
-            return response_text
+        return (response_texts, scores) if return_scores else response_texts
+
+    # ----------------- 同步包装接口 -----------------
+    def generate(self, input_list: List, batch_size=None, return_scores=False, **params) -> List[str]:
+        loop = get_background_loop()
+        future = asyncio.run_coroutine_threadsafe(
+            self._generate_async(input_list, batch_size=batch_size, return_scores=return_scores, **params),
+            loop
+        )
+        return future.result()

flashrag_dev-0.1.4.dev20250209/flashrag/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.4dev20250209"

{flashrag_dev-0.1.4.dev20250208 → flashrag_dev-0.1.4.dev20250209}/flashrag_dev.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.4.dev20250208
+Version: 0.1.4.dev20250209
 Summary: A library for efficient Retrieval-Augmented Generation research
 Home-page: https://github.com/RUC-NLPIR/FlashRAG
 Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou

{flashrag_dev-0.1.4.dev20250208 → flashrag_dev-0.1.4.dev20250209}/flashrag_dev.egg-info/requires.txt

@@ -18,6 +18,8 @@ fschat
 streamlit
 chonkie>=0.4.0
 gradio==5.11.0
+rouge-chinese
+jieba
 
 [core]
 datasets
@@ -40,6 +42,8 @@ fschat
 streamlit
 chonkie>=0.4.0
 gradio==5.11.0
+rouge-chinese
+jieba
 
 [full]
 datasets
@@ -62,6 +66,8 @@ fschat
 streamlit
 chonkie>=0.4.0
 gradio==5.11.0
+rouge-chinese
+jieba
 pyserini
 sentence-transformers>=3.0.1
 vllm==0.5.5

flashrag_dev-0.1.4.dev20250208/flashrag/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.4dev20250208"