26 lines
1.5 KiB
Python
26 lines
1.5 KiB
Python
import json
|
|
|
|
from django.core.management.base import BaseCommand, CommandError
|
|
|
|
from control_plane.model_studio.models import DatasetCurationProposal
|
|
from control_plane.model_studio.services import ModelStudioService
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Materialize a Qwen curation proposal into a new immutable Spark dataset version and source-hash splits."
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument("--proposal", required=True)
|
|
parser.add_argument("--output-directory", required=True)
|
|
parser.add_argument("--ssh-alias", default="spark")
|
|
parser.add_argument("--strict-schema-repair", action="store_true")
|
|
|
|
def handle(self, *args, **options):
|
|
proposal = DatasetCurationProposal.objects.filter(id=options["proposal"]).first()
|
|
if proposal is None:
|
|
raise CommandError("DatasetCurationProposal not found.")
|
|
training_project = proposal.source_versions.first().dataset.training_project if proposal.source_versions.exists() else None
|
|
if training_project is None:
|
|
raise CommandError("Proposal has no source versions.")
|
|
version = ModelStudioService().materialize_spark_curation(training_project, proposal, output_directory=options["output_directory"], ssh_alias=options["ssh_alias"], strict=options["strict_schema_repair"])
|
|
self.stdout.write(json.dumps({"dataset_version": str(version.id), "reference": version.manifest_reference, "records": version.record_count, "splits": version.split_metadata}, indent=2))
|