mirror of
https://github.com/mindverse/Second-Me.git
synced 2026-07-23 08:03:28 +00:00
91 lines
3.5 KiB
Python
91 lines
3.5 KiB
Python
from abc import ABC, abstractmethod
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import List, Dict, Any
|
|
|
|
from openai import OpenAI
|
|
from tqdm import tqdm
|
|
|
|
from lpm_kernel.api.services.user_llm_config_service import UserLLMConfigService
|
|
from lpm_kernel.configs.logging import get_train_process_logger
|
|
|
|
logger = get_train_process_logger()
|
|
|
|
|
|
class BaseData(ABC):
|
|
def __init__(self, is_cot: bool = True, max_workers: int = 10):
|
|
self.max_workers = max_workers
|
|
self.is_cot = is_cot
|
|
|
|
self.user_llm_config_service = UserLLMConfigService()
|
|
self.user_llm_config = self.user_llm_config_service.get_available_llm()
|
|
|
|
if self.user_llm_config is None:
|
|
self.client = None
|
|
self.model_name = None
|
|
self.reasoning_client = None
|
|
self.reasoning_model_name = None
|
|
else:
|
|
self.client = OpenAI(
|
|
api_key=self.user_llm_config.chat_api_key,
|
|
base_url=self.user_llm_config.chat_endpoint,
|
|
)
|
|
self.model_name = self.user_llm_config.chat_model_name
|
|
self.reasoning_client = OpenAI(
|
|
api_key=self.user_llm_config.thinking_api_key,
|
|
base_url=self.user_llm_config.thinking_endpoint,
|
|
)
|
|
self.reasoning_model_name = self.user_llm_config.thinking_model_name
|
|
|
|
def preprocess(self):
|
|
pass
|
|
|
|
def build_messages(self):
|
|
pass
|
|
|
|
def build_responses(self, messages_list: List[List[Dict[str, Any]]]):
|
|
|
|
def process_request(messages):
|
|
try:
|
|
if self.is_cot:
|
|
response = self.reasoning_client.chat.completions.create(
|
|
model=self.reasoning_model_name,
|
|
messages=messages,
|
|
extra_body={"metadata": {"tags": ["lpm-pipeline-using"]}},
|
|
)
|
|
result = response.choices[0].message
|
|
logger.info("===============================================")
|
|
logger.info(messages[1]["content"])
|
|
logger.info("-----------------------------------------------")
|
|
logger.info("<think>" + result.reasoning_content + "</think>")
|
|
logger.info("<answer>" + result.content + "</answer>")
|
|
logger.info("")
|
|
return "<think>" + result.reasoning_content + "</think>\n\n<answer>" + result.content + "</answer>"
|
|
else:
|
|
responses = self.client.chat.completions.create(
|
|
model=self.model_name,
|
|
messages=messages,
|
|
temperature=0,
|
|
max_tokens=3000,
|
|
extra_body={"metadata": {"tags": ["lpm-pipeline-using"]}},
|
|
).choices[0].message.content
|
|
return responses
|
|
except Exception as e:
|
|
logger.error(f"Raise ERROR: {e} WHEN GENERATE RESPONSE")
|
|
raise
|
|
|
|
with ThreadPoolExecutor(max_workers=max(min(self.max_workers, len(messages_list)), 1)) as executor:
|
|
futures = [executor.submit(process_request, messages) for messages in messages_list]
|
|
results = []
|
|
|
|
for future in tqdm(futures, total=len(messages_list), desc="Generating responses"):
|
|
result = future.result()
|
|
results.append(result)
|
|
|
|
return results
|
|
|
|
@abstractmethod
|
|
def run(self):
|
|
pass
|
|
|
|
def postprocess(self):
|
|
pass
|