diff --git a/build-dataset/config.yml b/build-dataset/config.yml index f4d342f..e740778 100644 --- a/build-dataset/config.yml +++ b/build-dataset/config.yml @@ -1,6 +1,6 @@ llm: provider: codex # openai, codex - # model: gpt-6-Astra + model: gpt-6-astra # model: qwen/qwen3-8b generation: @@ -8,7 +8,7 @@ generation: max_tokens: 4096 runtime: - concurrency: 1 # Колличество одновременных запросов + concurrency: 5 # Колличество одновременных запросов prompts: instructions: | diff --git a/build-dataset/robot-move-data.py b/build-dataset/robot-move-data.py index 0a9ce49..641bbb2 100644 --- a/build-dataset/robot-move-data.py +++ b/build-dataset/robot-move-data.py @@ -1,7 +1,5 @@ import asyncio -from pathlib import Path -import pandas as pd from dotenv import load_dotenv # from pydantic import BaseModel @@ -9,61 +7,56 @@ from pydantic_ai import Agent from pydantic_ai.settings import ModelSettings from llm.config import load_config -from llm.factory import LLMModelProvider -from llm.schemas import DatasetRecord, Message, RobotPlan, SFTFormat +from llm.factory import LLMModelProvider, DatasetWriter +from llm.schemas import (RobotPlan, RobotDatasetRecord, + Prompts, ModelResponse, + ModelInfo, GenerationSettings) load_dotenv() -def save_to_csv(record: DatasetRecord, - path: str | Path): - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - - exists = path.exists() - - df = pd.DataFrame([record.model_dump()]) - df.to_csv(path, - mode="a" if exists else "w", - header=not exists, - index=False, - encoding="utf-8") - -def save_to_jsonl(sft_format: SFTFormat, - path: str | Path): - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - - with path.open("a", encoding="utf-8") as file: - file.write(sft_format.model_dump_json()) - file.write("\n") - - -def record_to_sft(record: DatasetRecord) -> SFTFormat: - return SFTFormat( - message=[ - Message( - role="system", - content=record.system_prompt - ), - Message( - role="user", - content=record.user_prompt - ), - Message( - role="assistant", - content=record.assistant_content - ) - ] - ) - def format_robot_plan(plan: RobotPlan): return "\n".join( - f"{index}. {step}" - for index, step in enumerate(plan.steps, start=1) + f"{step}" + for step in plan.steps ) +async def generate_record(task: str, + agent: Agent, + settings: ModelSettings, + config, + semaphore: asyncio.Semaphore) -> RobotDatasetRecord: + + async with semaphore: + user_prompt = config.prompts.user.format(task=task) + + result = await agent.run(user_prompt, + model_settings=settings) + assistant_content = format_robot_plan(result.output) + + return RobotDatasetRecord( + task=task, + prompts=Prompts( + system_prompt=config.prompts.system, + instructions=config.prompts.instructions, + user_prompt=user_prompt + ), + response=ModelResponse( + model_response=result.output.model_dump_json(), + assistant_content=assistant_content + ), + model_info=ModelInfo( + provider=config.llm.provider, + model=config.llm.model + ), + generation=GenerationSettings( + temperature=config.generation.temperature, + max_tokens=config.generation.max_tokens + ) + ) + + async def main(): config = load_config("./build-dataset/config.yml") model = await LLMModelProvider.create_from_config(config.llm) @@ -81,30 +74,42 @@ async def main(): **config.generation.model_dump() ) - task="Перемести деталь A из точки PICK_A в точку PLACE_A. Безопасная высота подъема — 200 мм." - user_prompt = config.prompts.user.format(task=task) - result = await agent.run(user_prompt, - model_settings=settings) + tasks = [ + ( + "Перемести деталь A из точки PICK_A в точку PLACE_A. " + "Безопасная высота подъема — 200 мм." + ), + ( + "Перемести деталь B из точки PICK_B в точку PLACE_B. " + "Безопасная высота подъема — 150 мм." + ), + ( + "Возьми крышку из точки COVER_PICK и перемести " + "ее в точку COVER_PLACE." + ), + ] - assistant_content = format_robot_plan(result.output) - record = DatasetRecord( - task=task, - system_prompt=config.prompts.system, - instructions=config.prompts.instructions, - user_prompt=user_prompt, - model_response=result.output.model_dump_json(), - assistant_content=assistant_content, - provider=config.llm.provider, - model=config.llm.model, - temperature=config.generation.temperature, - max_tokens=config.generation.max_tokens - ) + semaphore = asyncio.Semaphore(config.runtime.concurrency) - save_to_csv(record=record, - path="./dataset/robot-move.csv") - sft_format = record_to_sft(record) - save_to_jsonl(sft_format=sft_format, - path="./dataset/robot-move.jsonl") + records = await asyncio.gather(*[ + generate_record( + task, + agent, + settings, + config, + semaphore, + ) + for task in tasks + ]) + + for record in records: + dataset = DatasetWriter( + dataset_record=record, + base_path="./dataset" + ) + + dataset.to_csv(filename="robot-plan.csv") + dataset.to_jsonl(filename="robotplan.jsonl") if __name__ == "__main__": diff --git a/llm/credentials.py b/llm/credentials.py index 0f141b0..d9e495a 100644 --- a/llm/credentials.py +++ b/llm/credentials.py @@ -10,12 +10,12 @@ from pydantic_ai.providers.openai_codex import ( class FileCodexCredentionalSource(OpenAICodexCredentialSource): def __init__(self, - path: Path, + path: str| Path, encoding: str="utf-8", ensure_ascii: bool=False, indent: int=2): super().__init__() - self._path = path + self._path = Path(path) self._encoding=encoding self._ensure_ascii=ensure_ascii self._indent=indent @@ -32,6 +32,8 @@ class FileCodexCredentionalSource(OpenAICodexCredentialSource): ) async def save(self, credentials: OpenAICodexCredentials) -> None: + self.path.parent.mkdir(parents=True , exist_ok=True) + self._path.write_text( json.dumps( asdict(credentials), diff --git a/llm/factory.py b/llm/factory.py index c21cd1d..b79cc9f 100644 --- a/llm/factory.py +++ b/llm/factory.py @@ -1,5 +1,6 @@ import os import webbrowser +import pandas as pd from pathlib import Path from pydantic_ai.models.openai import OpenAIChatModel @@ -12,6 +13,7 @@ from pydantic_ai.providers.openai_codex import ( from .config import CodexConfig, LLMConfig, OpenAiConfig from .credentials import FileCodexCredentionalSource +from .schemas import BaseDatasetRecord, ChatRecord, ChatMessage class LLMModelProvider: @@ -81,4 +83,69 @@ class LLMModelProvider: base_url=base_url ) - raise ValueError(f"Unsupported LLM config: {type(config)}") \ No newline at end of file + raise ValueError(f"Unsupported LLM config: {type(config)}") + +class DatasetWriter: + def __init__(self, + dataset_record: BaseDatasetRecord, + base_path: str | Path): + + + self._dataset_record = dataset_record + self._path = Path(base_path) + self._path.mkdir(parents=True, exist_ok=True) + + def __check_file_exists(self, filename :str)-> bool: + file_path = self._path / filename + return file_path.exists() + + def to_csv(self, + filename: str) -> None: + + exists=self.__check_file_exists(filename) + + df = pd.json_normalize( + self._dataset_record.model_dump() + ) + + df.to_csv(self._path / filename, + mode="a" if exists else "w", + header=not exists, + index=False, + encoding="utf-8") + + def to_chat_record(self) -> ChatRecord: + response = self._dataset_record.response + prompts = self._dataset_record.prompts + + assistant_content = ( + response.assistant_content + or response.model_response + ) + + return ChatRecord( + messages=[ + ChatMessage( + role="system", + content=prompts.system_prompt, + ), + ChatMessage( + role="user", + content=prompts.user_prompt, + ), + ChatMessage( + role="assistant", + content=assistant_content, + ), + ] + ) + + def to_jsonl(self, + filename: str) -> None: + + chat_record = self.to_chat_record() + path = self._path/filename + + with path.open("a", encoding='utf-8') as file: + file.write(chat_record.model_dump_json()) + file.write("\n") \ No newline at end of file diff --git a/llm/schemas.py b/llm/schemas.py index e27bcad..a7a90b4 100644 --- a/llm/schemas.py +++ b/llm/schemas.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Literal, Optional from pydantic import BaseModel, Field @@ -10,7 +10,7 @@ class RobotPlan(BaseModel): ) -class Message(BaseModel): +class ChatMessage(BaseModel): role: Literal["system", "user", "assistant"] @@ -18,22 +18,32 @@ class Message(BaseModel): content: str -class SFTFormat(BaseModel): - message: list[Message] - - -class DatasetRecord(BaseModel): - task: str +class ChatRecord(BaseModel): + messages: list[ChatMessage] +class Prompts(BaseModel): system_prompt: str instructions: str user_prompt: str - model_response: str - assistant_content: str +class ModelResponse(BaseModel): + model_response: str + assistant_content: Optional[str] = None + +class ModelInfo(BaseModel): provider: str model: str +class GenerationSettings(BaseModel): temperature: float - max_tokens: int \ No newline at end of file + max_tokens: int + +class BaseDatasetRecord(BaseModel): + prompts: Prompts + response: ModelResponse + model_info: ModelInfo + generation: GenerationSettings + +class RobotDatasetRecord(BaseDatasetRecord): + task: str