Merge pull request #14 from Open-DataFlow/sunnyhaze

Revise calling style to nn.module style

Author: SunnyHaze

.gitignore

@@ -4,4 +4,8 @@
 *.egg
 *.egg-info  
 /dataflow/example/ReasoningPipeline/pipeline_math_step*.jsonl
-!/dataflow/example/ReasoningPipeline/pipeline_math.json
+!/dataflow/example/ReasoningPipeline/pipeline_math.json
+
+test/example
+
+cache

dataflow/__init__.py

@@ -1,8 +1,12 @@
 from .utils import *
 from .version import __version__, version_info
+from .logger import get_logger
 
-
-__all__ = ['__version__', 'version_info']
+__all__ = [
+    '__version__',
+    'version_info',
+    'get_logger',
+]
 
 
 

dataflow/utils/Generator.py dataflow/core/Generator.pydataflow/utils/Generator.py renamed to dataflow/core/Generator.py

@@ -2,7 +2,9 @@
 from typing import Any, List
 
 
-class Generator(ABC):
+class GeneratorABC(ABC):
+    """Abstract base class for data generators. Which may be used to generate data from a model or API. Called by operators
+    """
     @abstractmethod
     def generate(self) -> Any:
         """
@@ -12,7 +14,7 @@ def generate(self) -> Any:
         pass
 
     @abstractmethod
-    def generate_from_input(self, input: List[str]) -> List[str]:
+    def generate_from_input(self, input: List[str], system_prompt: str) -> List[str]:
         """
         Generate data from input.
         input: List[str], the input of the generator

dataflow/core/Operator.py

@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+from dataflow.logger import get_logger
+
+class OperatorABC(ABC):
+
+    # @abstractmethod
+    # def check_config(self, config: dict) -> None:
+    #     """
+    #     Check the config of the operator. If config lacks any required keys, raise an error.
+    #     """
+    #     pass
+    
+    @abstractmethod
+    def run(self) -> None:
+        """
+        Main function to run the operator.
+        """
+        pass
+
+def get_operator(operator_name, args) -> OperatorABC:
+    from dataflow.utils import OPERATOR_REGISTRY
+    print(operator_name, args)
+    operator = OPERATOR_REGISTRY.get(operator_name)(args)
+    logger = get_logger()
+    if operator is not None:
+        logger.info(f"Successfully get operator {operator_name}, args {args}")
+    else:
+        logger.error(f"operator {operator_name} is not found")
+    assert operator is not None
+    print(operator)
+    return operator

dataflow/core/__init__.py

@@ -0,0 +1,7 @@
+from .Operator import OperatorABC, get_operator
+from .Generator import GeneratorABC
+__all__ = [
+    'OperatorABC',
+    'get_operator',
+    'GeneratorABC',    
+]

dataflow/example/ReasoningPipeline/pipeline_math_short.json

@@ -2,11 +2,11 @@
 import aisuite as ai
 import pandas as pd
 from tqdm import tqdm
-from dataflow.utils.Generator import Generator
+from dataflow.core import GeneratorABC
 from dataflow.utils.Storage import FileStorage
 from typing import List
 
-class APIGenerator_aisuite(Generator):
+class APIGenerator_aisuite(GeneratorABC):
     def __init__(self, config: dict):
         configs = config  # Assuming config.configs is a list of configurations
 

dataflow/utils/APIGenerator_aisuite.py …aflow/generators/APIGenerator_aisuite.pydataflow/utils/APIGenerator_aisuite.py renamed to dataflow/generators/APIGenerator_aisuite.py dataflow/utils/APIGenerator_aisuite.py renamed to dataflow/generators/APIGenerator_aisuite.py

@@ -2,32 +2,37 @@
 import requests
 import os
 import logging
-import pandas as pd
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from tqdm import tqdm
-from dataflow.utils.Storage import FileStorage
-from dataflow.utils.Generator import Generator
+from dataflow.core  import GeneratorABC
 import re
 
-class APIGenerator_request(Generator):
-    def __init__(self, config: dict):
-        self.config = config
-        
+class APIGenerator_request(GeneratorABC):
+    """Use OpenAI API to generate responses based on input messages.
+    """
+    def __init__(self, 
+                 api_url: str = "https://api.openai.com/v1/chat/completions",
+                 model_name: str = "gpt-4o",
+                 max_workers: int = 10
+                 ):
         # Get API key from environment variable or config
-        self.api_url = self.config.get("api_url", "https://api.openai.com/v1/chat/completions")
+        self.api_url = api_url
+        self.model_name = model_name
+        self.max_workers = max_workers
+
+        # config api_key in os.environ global, since safty issue.
         self.api_key = os.environ.get("API_KEY")
         if self.api_key is None:
             raise ValueError("Lack of API_KEY")
 
-        self.datastorage = FileStorage(self.config)
-    
+    """corden due to I don't confindent about implementation——Sunnyhaze
     def check_config(self):
         # Ensure all necessary keys are in the config
         necessary_keys = ['input_file', 'output_file', 'input_key', 'output_key', 'max_workers']
         for key in necessary_keys:
             if key not in self.config:
                 raise ValueError(f"Key {key} is missing in the config")
-            
+    """    
     def format_response(self, response: dict) -> str:    
         # check if content is formatted like <think>...</think>...<answer>...</answer>
         content = response['choices'][0]['message']['content']
@@ -74,6 +79,9 @@ def api_chat(self, system_info: str, messages: str, model: str):
             logging.error(f"API request error: {e}")
             return None
 
+    def generate(self):
+        pass # for develop, TODO
+    """ Corden due to I don't confindent about implementation——Sunnyhaze
     def generate(self):
         self.check_config()
         # Read input file
@@ -105,8 +113,8 @@ def generate(self):
         raw_dataframe[self.config['output_key']] = responses
         self.datastorage.write(self.config['output_file'], raw_dataframe)
         return
-
-    def generate_from_input(self, input: list[str]) -> list[str]:
+    """
+    def generate_from_input(self, input: list[str], system_prompt: str = "") -> list[str]:
         def api_chat_with_id(system_info: str, messages: str, model: str, id):
             try:
                 payload = json.dumps({
@@ -136,17 +144,18 @@ def api_chat_with_id(system_info: str, messages: str, model: str, id):
                 logging.error(f"API request error: {e}")
                 return id,None
         responses = [None] * len(input)
-        
+        # -- end of subfunction api_chat_with_id --
+
         # 使用 ThreadPoolExecutor 并行处理多个问题
         # logging.info(f"Generating {len(questions)} responses")
-        with ThreadPoolExecutor(max_workers=self.config['max_workers']) as executor:
+        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
             futures = [
                 executor.submit(
                     api_chat_with_id,
-                    self.config['system_prompt'],
-                    question,
-                    self.config['model_name'],
-                    idx
+                    system_info = system_prompt,
+                    messages = question,
+                    model = self.model_name,
+                    id = idx
                 ) for idx, question in enumerate(input)
             ]
             for future in tqdm(as_completed(futures), total=len(futures), desc="Generating......"):

dataflow/utils/APIGenerator_request.py …aflow/generators/APIGenerator_request.pydataflow/utils/APIGenerator_request.py renamed to dataflow/generators/APIGenerator_request.py dataflow/utils/APIGenerator_request.py renamed to dataflow/generators/APIGenerator_request.py

@@ -3,9 +3,9 @@
 from huggingface_hub import snapshot_download
 import pandas as pd
 from dataflow.utils.Storage import FileStorage
-from dataflow.utils.Generator import Generator
+from dataflow.core import GeneratorABC
 
-class LocalModelGenerator(Generator):
+class LocalModelGenerator(GeneratorABC):
     '''
     A class for generating text using vllm, with model from huggingface or local directory
     '''

dataflow/utils/LocalModelGenerator.py …taflow/generators/LocalModelGenerator.pydataflow/utils/LocalModelGenerator.py renamed to dataflow/generators/LocalModelGenerator.py dataflow/utils/LocalModelGenerator.py renamed to dataflow/generators/LocalModelGenerator.py

@@ -0,0 +1,9 @@
+from .APIGenerator_aisuite import APIGenerator_aisuite
+from .APIGenerator_request import APIGenerator_request
+from .LocalModelGenerator import LocalModelGenerator
+
+__all__ = [
+    "APIGenerator_aisuite",
+    "APIGenerator_request",
+    "LocalModelGenerator"
+]