116 lines
3.5 KiB
Python
116 lines
3.5 KiB
Python
import asyncio
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
# from pydantic import BaseModel
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.settings import ModelSettings
|
|
|
|
from llm.config import load_config
|
|
from llm.factory import LLMModelProvider, DatasetWriter
|
|
from llm.schemas import (RobotPlan, RobotDatasetRecord,
|
|
Prompts, ModelResponse,
|
|
ModelInfo, GenerationSettings)
|
|
|
|
load_dotenv()
|
|
|
|
|
|
def format_robot_plan(plan: RobotPlan):
|
|
return "\n".join(
|
|
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)
|
|
|
|
|
|
# model = await LLMModelProvider.create_codex(model_name='gpt-5.6-luna')
|
|
# model = LLMModelProvider.create_openai_compatible(model_name="qwen/qwen3-8b",
|
|
# api_key=os.getenv("API_KEY"),
|
|
# base_url=os.getenv("BASE_URL"))
|
|
agent = Agent(model=model,
|
|
output_type=RobotPlan,
|
|
system_prompt=config.prompts.system,
|
|
instructions=config.prompts.instructions)
|
|
settings = ModelSettings(
|
|
**config.generation.model_dump()
|
|
)
|
|
|
|
tasks = [
|
|
(
|
|
"Перемести деталь A из точки PICK_A в точку PLACE_A. "
|
|
"Безопасная высота подъема — 200 мм."
|
|
),
|
|
(
|
|
"Перемести деталь B из точки PICK_B в точку PLACE_B. "
|
|
"Безопасная высота подъема — 150 мм."
|
|
),
|
|
(
|
|
"Возьми крышку из точки COVER_PICK и перемести "
|
|
"ее в точку COVER_PLACE."
|
|
),
|
|
]
|
|
|
|
semaphore = asyncio.Semaphore(config.runtime.concurrency)
|
|
|
|
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__":
|
|
asyncio.run(main()) |