wehub-resource-sync
380 lines
15 KiB
Python
380 lines
15 KiB
Python
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import subprocess
|
|
from dataclasses import dataclass
|
|
from itertools import cycle
|
|
from pathlib import Path
|
|
from typing import Dict, List
|
|
|
|
import aiofiles
|
|
from openai import AsyncOpenAI
|
|
from tenacity import retry, stop_after_attempt, wait_random_exponential
|
|
from tqdm.asyncio import tqdm
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class RangeSet:
|
|
"""Manage processed line ranges with efficient storage and querying"""
|
|
|
|
ranges: List[tuple]
|
|
|
|
def add(self, number: int):
|
|
"""Add a number to the range set and merge adjacent ranges"""
|
|
new_ranges = []
|
|
added = False
|
|
for start, end in sorted(self.ranges):
|
|
if number < start - 1:
|
|
if not added:
|
|
new_ranges.append((number, number))
|
|
added = True
|
|
new_ranges.append((start, end))
|
|
elif number == start - 1:
|
|
new_ranges.append((number, end))
|
|
added = True
|
|
elif number <= end:
|
|
new_ranges.append((start, end))
|
|
added = True
|
|
else:
|
|
new_ranges.append((start, end))
|
|
if not added:
|
|
new_ranges.append((number, number))
|
|
self.ranges = self.merge_ranges(new_ranges)
|
|
|
|
@staticmethod
|
|
def merge_ranges(ranges: List[tuple]) -> List[tuple]:
|
|
"""Merge overlapping or adjacent ranges"""
|
|
if not ranges:
|
|
return []
|
|
sorted_ranges = sorted(ranges)
|
|
merged = [sorted_ranges[0]]
|
|
for current in sorted_ranges[1:]:
|
|
last = merged[-1]
|
|
if current[0] <= last[1] + 1:
|
|
merged[-1] = (last[0], max(last[1], current[1]))
|
|
else:
|
|
merged.append(current)
|
|
return merged
|
|
|
|
def contains(self, number: int) -> bool:
|
|
"""Check if a number exists in any range"""
|
|
for start, end in self.ranges:
|
|
if start <= number <= end:
|
|
return True
|
|
return False
|
|
|
|
def to_file_format(self) -> str:
|
|
"""Serialize ranges to compact string format"""
|
|
return ",".join(f"{start}-{end}" if start != end else str(start) for start, end in self.ranges)
|
|
|
|
@classmethod
|
|
def from_file(cls, content: str) -> "RangeSet":
|
|
"""Deserialize from string format"""
|
|
if not content:
|
|
return cls(ranges=[])
|
|
ranges = []
|
|
for part in content.split(","):
|
|
if "-" in part:
|
|
start, end = map(int, part.split("-"))
|
|
ranges.append((start, end))
|
|
else:
|
|
num = int(part)
|
|
ranges.append((num, num))
|
|
return cls(ranges=ranges)
|
|
|
|
@property
|
|
def processed_count(self) -> int:
|
|
"""Total number of processed items"""
|
|
return sum(end - start + 1 for start, end in self.ranges)
|
|
|
|
|
|
class OpenAIClientPool:
|
|
"""Manage round-robin distribution of API clients"""
|
|
|
|
def __init__(self, base_urls: List[str], api_keys: List[str] = None):
|
|
if isinstance(base_urls, str):
|
|
base_urls = base_urls.split(",")
|
|
if isinstance(api_keys, str):
|
|
api_keys = api_keys.split(",")
|
|
|
|
if api_keys is None:
|
|
api_keys = ["NONE" for _ in range(len(base_urls))]
|
|
|
|
if len(api_keys) != len(base_urls):
|
|
raise ValueError("API keys and base URLs should have the same length!")
|
|
|
|
self.clients = cycle([AsyncOpenAI(base_url=url, api_key=key) for url, key in zip(base_urls, api_keys)])
|
|
|
|
def get_client(self) -> AsyncOpenAI:
|
|
"""Get next available client in rotation"""
|
|
return next(self.clients)
|
|
|
|
|
|
class OpenAIProcessor:
|
|
"""Async processor for batch processing with OpenAI-compatible APIs"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_file: str,
|
|
output_file: str,
|
|
prompt_key: str,
|
|
base_urls: List[str],
|
|
api_keys: List[str] = None,
|
|
prompt_suffix: str = "",
|
|
status_file: str = "status.txt",
|
|
concurrency: int = 8,
|
|
model: str = "deepseek-r1",
|
|
temperature: float = 0.6,
|
|
top_p: float = 1.0,
|
|
max_tokens: int = 65536,
|
|
timeout: int = 3600,
|
|
response_key: str = "response",
|
|
reasoning_key: str = "reasoning",
|
|
):
|
|
self.input_file = input_file
|
|
self.output_file = output_file
|
|
self.client_pool = OpenAIClientPool(base_urls, api_keys)
|
|
self.status_file = status_file
|
|
self.processed_set = RangeSet([])
|
|
self.concurrency = concurrency
|
|
self.model = model
|
|
self.temperature = temperature
|
|
self.top_p = top_p
|
|
self.max_tokens = max_tokens
|
|
self.total_lines = 0
|
|
self.progress_bar = None
|
|
self.write_lock = asyncio.Lock()
|
|
self.status_lock = asyncio.Lock()
|
|
self.prompt_key = prompt_key
|
|
self.prompt_suffix = prompt_suffix
|
|
self.timeout = timeout
|
|
self.response_key = response_key
|
|
self.reasoning_key = reasoning_key
|
|
|
|
self._load_status()
|
|
|
|
def _load_status(self):
|
|
"""Load processing status from file"""
|
|
try:
|
|
with open(self.status_file, "r", encoding="utf-8") as f:
|
|
content = f.read().strip()
|
|
self.processed_set = RangeSet.from_file(content)
|
|
logger.info(f"Resumed processed ranges: {self.processed_set.to_file_format()}")
|
|
except FileNotFoundError:
|
|
self.processed_set = RangeSet([])
|
|
|
|
async def _save_status(self):
|
|
"""Save current processing status to file"""
|
|
async with self.status_lock:
|
|
content = self.processed_set.to_file_format()
|
|
async with aiofiles.open(self.status_file, "w", encoding="utf-8") as f:
|
|
await f.write(content)
|
|
|
|
def _count_total_lines(self) -> int:
|
|
"""Count total lines in input file"""
|
|
try:
|
|
result = subprocess.run(
|
|
["wc", "-l", self.input_file], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
|
|
)
|
|
return int(result.stdout.strip().split()[0])
|
|
except Exception as e:
|
|
print(f"Failed to count lines using `wc -l` command: {str(e)}")
|
|
with open(self.input_file, "r", encoding="utf-8") as f:
|
|
return sum(1 for _ in f)
|
|
|
|
async def _line_generator(self):
|
|
"""Generate unprocessed lines with line numbers"""
|
|
self.total_lines = self._count_total_lines()
|
|
with open(self.input_file, "r", encoding="utf-8") as f:
|
|
for line_num, line in enumerate(f):
|
|
if not self.processed_set.contains(line_num):
|
|
yield line_num, json.loads(line.strip())
|
|
|
|
@retry(stop=stop_after_attempt(5), wait=wait_random_exponential(multiplier=1, max=60))
|
|
async def _call_openai(self, client: AsyncOpenAI, line_num: int, data: Dict[str, str]) -> str:
|
|
"""Execute API call with retry logic"""
|
|
try:
|
|
# Prepend processing instructions to the content
|
|
content = data.get(self.prompt_key, "")
|
|
if isinstance(content, (tuple, list)):
|
|
content = content[0]
|
|
|
|
response = await client.chat.completions.create(
|
|
model=self.model,
|
|
messages=[{"role": "user", "content": content + self.prompt_suffix}],
|
|
temperature=self.temperature,
|
|
top_p=self.top_p,
|
|
max_tokens=self.max_tokens,
|
|
timeout=self.timeout,
|
|
stream=False,
|
|
)
|
|
return {"line_num": line_num, **data, **self._parse_response(response)}
|
|
|
|
except Exception as e:
|
|
logger.error(f"API call failed: {str(e)}")
|
|
if "rate limit" in str(e).lower():
|
|
await asyncio.sleep(5)
|
|
raise
|
|
|
|
def _parse_response(self, response) -> Dict[str, str]:
|
|
"""Parse API response into structured format"""
|
|
response_text = response.result
|
|
reasoning_text = ""
|
|
|
|
if "</think>" in response_text and response_text.count("</think>") == 1:
|
|
reasoning_text, _, response_text = response_text.partition("</think>")
|
|
|
|
if reasoning_text and not reasoning_text.startswith("<think>"):
|
|
reasoning_text = f"<think>\n{reasoning_text.strip()}"
|
|
if reasoning_text and not reasoning_text.endswith("</think>"):
|
|
reasoning_text = f"{reasoning_text.strip()}\n</think>"
|
|
|
|
return {
|
|
self.response_key: response_text,
|
|
self.reasoning_key: reasoning_text,
|
|
}
|
|
|
|
async def _write_result(self, line_num: int, result: Dict[str, str]):
|
|
"""Write processed result and update status"""
|
|
async with self.write_lock:
|
|
# Append result to output file
|
|
async with aiofiles.open(self.output_file, "a", encoding="utf-8") as f:
|
|
await f.write(json.dumps(result, ensure_ascii=False) + "\n")
|
|
|
|
# Update processing status
|
|
self.processed_set.add(line_num)
|
|
await self._save_status()
|
|
|
|
async def worker(self, queue: asyncio.Queue):
|
|
"""Process items from the queue"""
|
|
client = self.client_pool.get_client()
|
|
while True:
|
|
line_num, data = await queue.get()
|
|
try:
|
|
result = await self._call_openai(client, line_num, data)
|
|
await self._write_result(line_num, result)
|
|
self.progress_bar.update(1)
|
|
except Exception as e:
|
|
logger.error(f"Failed to process line {line_num}: {str(e)}")
|
|
finally:
|
|
queue.task_done()
|
|
|
|
async def run(self):
|
|
"""Main processing loop"""
|
|
total_lines = self._count_total_lines()
|
|
remaining = total_lines - self.processed_set.processed_count
|
|
|
|
if remaining <= 0:
|
|
logger.info("No data requires distilling!")
|
|
return
|
|
|
|
# Initialize progress bar with current progress
|
|
self.progress_bar = tqdm(
|
|
total=total_lines,
|
|
desc=f"[{self.model}] Data Distilling Progress",
|
|
dynamic_ncols=True,
|
|
initial=self.processed_set.processed_count,
|
|
)
|
|
|
|
queue = asyncio.Queue(maxsize=self.concurrency * 2)
|
|
|
|
# Ensure output file exists
|
|
async with aiofiles.open(self.output_file, "a", encoding="utf-8"):
|
|
pass
|
|
|
|
# Start worker tasks
|
|
workers = [asyncio.create_task(self.worker(queue)) for _ in range(self.concurrency)]
|
|
|
|
try:
|
|
# Feed unprocessed items to queue
|
|
async for line_num, data in self._line_generator():
|
|
await queue.put((line_num, data))
|
|
|
|
await queue.join()
|
|
finally:
|
|
# Cleanup resources
|
|
for worker_task in workers:
|
|
worker_task.cancel()
|
|
await asyncio.gather(*workers, return_exceptions=True)
|
|
self.progress_bar.close()
|
|
logger.info(f"Processing complete. Total processed: {self.processed_set.processed_count} / {total_lines}.")
|
|
logger.info("=" * 100)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="OpenAI Multi-Processing Interface")
|
|
parser.add_argument("--input_file", type=Path, required=True, help="Input JSONL filename")
|
|
parser.add_argument("--output_dir", type=Path, required=True, help="Output directory path")
|
|
parser.add_argument("--prompt_key", required=True, help="Key name for the input JSONL data prompt")
|
|
parser.add_argument("--response_key", required=True, help="Key name for the output JSONL data response")
|
|
parser.add_argument("--reasoning_key", required=True, help="Key name for the output JSONL data reasoning")
|
|
parser.add_argument("--base_urls", required=True, help="Comma-separated list of API endpoints")
|
|
parser.add_argument("--api_keys", default=None, help="Comma-separated list of API keys, Default: `None`")
|
|
parser.add_argument("--model", default="deepseek-r1", type=str, help="Model name to use, Default: `deepseek-r1`")
|
|
parser.add_argument("--prompt_suffix", default="", type=str, help="Suffix appended after each prompt, Default: ``")
|
|
parser.add_argument("--temperature", type=float, default=0.6, help="Sampling temperature, Default: `0.6`")
|
|
parser.add_argument("--top_p", type=float, default=1.0, help="Top-p sampling value, Default: `1.0`")
|
|
parser.add_argument("--timeout", type=int, default=3600, help="API request timeout in seconds, Default: `3600`s")
|
|
parser.add_argument(
|
|
"--max_tokens", type=int, default=65536, help="Maximum number of tokens to generate, Default: `65536`"
|
|
)
|
|
parser.add_argument(
|
|
"--concurrency", type=int, default=8, help="Maximum number of concurrent threads, Default: `8`"
|
|
)
|
|
parser.add_argument("--status_file", default=None, help="Status file path, Default: `None`")
|
|
parser.add_argument("--logging_file", default=None, help="Logging file path, Default: `None`")
|
|
args = parser.parse_args()
|
|
|
|
if not str(args.input_file).endswith(".jsonl"):
|
|
raise NotImplementedError("Currently only JSONL files are supported!")
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
filename_prefix = "distilled-" + args.input_file.stem
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s - %(levelname)s - %(message)s",
|
|
handlers=[
|
|
logging.FileHandler(args.logging_file or args.output_dir / f"{filename_prefix}.log"),
|
|
],
|
|
)
|
|
PROCESSOR_CONFIG = {
|
|
"input_file": args.input_file,
|
|
"output_file": args.output_dir / f"{filename_prefix}.jsonl",
|
|
"status_file": args.status_file or args.output_dir / f"{filename_prefix}.status",
|
|
"prompt_key": args.prompt_key,
|
|
"response_key": args.response_key,
|
|
"reasoning_key": args.reasoning_key,
|
|
"prompt_suffix": args.prompt_suffix,
|
|
"base_urls": args.base_urls,
|
|
"api_keys": args.api_keys,
|
|
"concurrency": args.concurrency,
|
|
"model": args.model,
|
|
"temperature": args.temperature,
|
|
"top_p": args.top_p,
|
|
"timeout": args.timeout,
|
|
"max_tokens": args.max_tokens,
|
|
}
|
|
processor = OpenAIProcessor(**PROCESSOR_CONFIG)
|
|
asyncio.run(processor.run())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|