111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
import asyncio
|
|
from pathlib import Path
|
|
|
|
import pandas as pd
|
|
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
|
|
from llm.schemas import DatasetRecord, Message, RobotPlan, SFTFormat
|
|
|
|
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)
|
|
)
|
|
|
|
|
|
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()
|
|
)
|
|
|
|
task="Перемести деталь A из точки PICK_A в точку PLACE_A. Безопасная высота подъема — 200 мм."
|
|
user_prompt = config.prompts.user.format(task=task)
|
|
result = await agent.run(user_prompt,
|
|
model_settings=settings)
|
|
|
|
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
|
|
)
|
|
|
|
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")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main()) |