Schema for DJL

Author: Jonathan Makunga

src/sagemaker/serve/builder/model_builder.py

@@ -609,17 +609,16 @@ def build(
             if self._is_jumpstart_model_id():
                 return self._build_for_jumpstart()
             if self._is_djl():  # pylint: disable=R1705
+                if self.schema_builder is None:
+                    self._schema_builder_init("text-generation")
+
                 return self._build_for_djl()
             else:
-                logger.info("******************************************************")
-                logger.info(f"schema_builder is None: {self.schema_builder is None}")
-
                 hf_model_md = get_huggingface_model_metadata(
                     self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
                 )
 
                 hf_task = hf_model_md.get("pipeline_tag")
-                logger.info(f"hf_task: {hf_task}")
                 if self.schema_builder is None:
                     self._schema_builder_init(hf_task)
 
@@ -688,15 +687,12 @@ def _schema_builder_init(self, model_task: str):
 
         try:
             sample_inputs, sample_outputs = task.retrieve_local_schemas(model_task)
-            logger.info(f"Sample input: {sample_inputs}")
-            logger.info(f"Sample output: {sample_outputs}")
         except ValueError:
             # TODO: try to retrieve schemas remotely
             pass
 
         if sample_inputs and sample_outputs:
             self.schema_builder = SchemaBuilder(sample_inputs, sample_outputs)
-            logger.info(f"schema_builder is not None: {self.schema_builder is None}")
         else:
             # TODO: Raise ClientError
             pass