LLAMA3.1 logging using mlflow

HemantvIkani32
New Contributor II

‎09-03-2025 07:59 AM

Hi team, I am facing issue after registering the meta llama 3.1 8b instruct model, while serving i get the response like

{error code: BAD REQUEST, "message": "Time out while evaluating the model. Verify that the model evaluate within the timeframe."

Given below my code snippet (IF there is any other way to register please help):

input_schema = Schema([ColSpec("string", "prompt"),
                       ColSpec("double", "temperature", required=False),
                       ColSpec("integer", "max_tokens", required=False),
                       ColSpec("string", "stop", required=False),
                       ColSpec("integer", "candidate_count", required=False)])

 
output_schema = Schema([ColSpec('string', 'predictions')])
signature = ModelSignature(inputs=input_schema, outputs=output_schema)
 input_example = {"prompt": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\nWhat is Apache Spark?\n\n### Response:\n"}
mlflow.set_registry_uri("databricks")

with mlflow.start_run() as run:
    result = mlflow.transformers.log_model(transformers_model={ "model": model,"tokenizer": tokenizer},
                                           task = "text-generation",
                                           artifact_path=REGISTERED_MODEL_NAME,
                                           registered_model_name=REGISTERED_MODEL_NAME,
                                           signature=signature,
                                          input_example=input_example,
                                           await_registration_for = 3600)
                                          


 

  

0 Kudos