Skip to content

Commit ca58746

Browse files
committed
Updated generate_inference_job to create ModelVersion if needed
1 parent 55c677b commit ca58746

File tree

1 file changed

+19
-1
lines changed

1 file changed

+19
-1
lines changed

sde_collections/models/collection.py

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -666,9 +666,27 @@ def generate_inference_job(self, classification_type):
666666

667667
InferenceJob = apps.get_model("inference", "InferenceJob")
668668
ModelVersion = apps.get_model("inference", "ModelVersion")
669+
670+
try:
671+
model_version = ModelVersion.get_active_version(classification_type)
672+
except ModelVersion.DoesNotExist:
673+
if classification_type == 1: # TDAMM
674+
model_name = "TDAMM"
675+
elif classification_type == 2: # DIVISION
676+
model_name = "DC"
677+
else:
678+
raise ValueError(f"Unsupported classification type: {classification_type}")
679+
680+
model_version = ModelVersion.objects.create(
681+
api_identifier=model_name,
682+
description=f"{model_name.upper()} classification model",
683+
classification_type=classification_type,
684+
is_active=True,
685+
)
686+
669687
return InferenceJob.objects.create(
670688
collection=self,
671-
model_version=ModelVersion.get_active_version(classification_type),
689+
model_version=model_version,
672690
)
673691

674692
def queue_necessary_classifications(self):

0 commit comments

Comments
 (0)