I am getting following error. I dont understand to resolve this my model is logged in mlflow. I am trying to load it locally in my colab.
Model Logged to MLFLOW code:
# Enable native 2x faster inference
FastLanguageModel.for_inference(model)
inference_pipeline = hf_pipeline("text-generation", model=model, tokenizer=tokenizer)
# Log the inference pipeline to MLflow
mlflow.transformers.log_model(
transformers_model=inference_pipeline,
artifact_path="lora_model_unsloth_PEFTllama3.2_3B",
registered_model_name="lora_model_unsloth_PEFTllama3.2_3B"
)
ML Load in Colab code:
import mlflow
# Set DagsHub as the tracking URI
mlflow.set_tracking_uri("https://dagshub.com/kushwanthkc/lora_model_unsloth_PEFTllama3.2_3B.mlflow")
# Use the full run ID as the model URI
logged_model = 'runs:/e73e9aba105c47d2bfdd58ef135f1478/lora_model_unsloth_PEFTllama3.2_3B' # Assuming this path is correct on DagsHub
# Load model as a PyFuncModel.
# Removing the 'device' argument as it's not supported by mlflow.pyfunc.load_model()
loaded_model = mlflow.pyfunc.load_model(logged_model)
error loaded_model:
ValueError: The model has been loaded with `accelerate` and therefore cannot be moved to a specific device. Please discard the `device` argument when creating your pipeline object.