Artifex/control_plane/model_studio/management/commands/modelstudio_dataset_curate.py
2026-08-17 02:04:43 +07:00

27 lines
1.5 KiB
Python

import json
from django.core.management.base import BaseCommand, CommandError
from control_plane.model_studio.models import DatasetVersion, TrainingProject
from control_plane.model_studio.services import ModelStudioService
from model_router.providers import providers_from_resources
from model_router.router import ModelRouter
class Command(BaseCommand):
help = "Ask the Qwen Dataset Curator for a structured, non-mutating Guard dataset improvement proposal."
def add_arguments(self, parser):
parser.add_argument("--project", required=True, help="TrainingProject slug")
parser.add_argument("--status", default="WARNING", help="Dataset validation status to review")
def handle(self, *args, **options):
project = TrainingProject.objects.filter(slug=options["project"]).first()
if project is None:
raise CommandError("TrainingProject not found.")
versions = list(DatasetVersion.objects.filter(dataset__training_project=project, validation_status=options["status"]).order_by("-record_count"))
if not versions:
raise CommandError("No matching DatasetVersions to curate.")
router = ModelRouter(providers_from_resources(), persist_requests=True)
proposal = ModelStudioService(router=router).propose_dataset_curation(project, versions)
self.stdout.write(json.dumps({"proposal_id": str(proposal.id), "title": proposal.title, "operations": proposal.proposed_operations, "validation_plan": proposal.validation_plan}, indent=2, default=str))