from sklearn import datasets

from sklearn.model_selection import train_test_split

from sklearn.ensemble import RandomForestClassifier

from sklearn.metrics import accuracy_score

from mlflow.models.signature import ModelSignature

from mlflow.types.schema import Schema, ColSpec

import mlflow

import mlflow.pyfunc

import pandas as pd

# Model wrapper class

class ModelWrapper(mlflow.pyfunc.PythonModel):

    # Initialize model in the constructor

    def __init__(self, model):

        self.model = model

    # Prediction function

    def predict(self, context, model_input):

        # Predict the probabilities and class

        prediction_probabilities = self.model.predict_proba(model_input)

        predictions = self.model.predict(model_input)

        # Create a DataFrame to hold the results

        result = pd.DataFrame(prediction_probabilities, columns=['prob_0', 'prob_1', 'prob_2'])

        # Inside the predict function:

        class_labels = ["setosa", "versicolor", "virginica"]

        predictions = self.model.predict(model_input)

        predicted_probabilities = self.model.predict_proba(model_input)

        result = pd.DataFrame(predicted_probabilities, columns=[f'prob_{label}' for label in class_labels])

        result['prediction'] = [class_labels[prediction] for prediction in predictions]

        return result

# Load iris dataset and split the dataset into train and test sets

iris = datasets.load_iris()

X = iris.data

y = iris.target

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# Initialize and train RandomForestClassifier

rf = RandomForestClassifier(random_state=42)

rf.fit(X_train, y_train)