AttributeError: module 'keras.api._v2.keras' has no attribute '__version__'

I am getting error something like this

AttributeError                            Traceback (most recent call last)
Cell In[10], line 10
      8     evaluation.log_into_mlflow()
      9 except Exception as e:
---> 10     raise e

Cell In[10], line 8
      6     evaluation.evaluation()
      7     evaluation.save_score()
----> 8     evaluation.log_into_mlflow()
      9 except Exception as e:
     10     raise e

Cell In[9], line 62
     57 # Model registry does not work with file store
     58 if tracking_url_type_store != "file":
     59 # Register the model with specified flavor options
     60     flavor_options = {
     61         "model_type": keras,
---> 62         "keras_version": tf.keras.__version__,
     63         "save_format": "h5"  
     64     }
     66     # Register the model
     67     # There are other ways to use the Model Registry, which depends on the use case,
     68     # please refer to the doc for more information:
     69     # https://mlflow.org/docs/latest/model-registry.html#api-workflow
     70     mlflow.keras.log_model(self.model, "model", registered_model_name="VGG16Model", keras_version=tf.keras.__version__,flavor=flavor_options)

I am building an MLOps project for Renal Health classification I am facing this issue in the model evaluation stage where I am using MLFlow for this tracing and I am getting this issue while pushing this to MLFlow tracking. Below is the code snippet what I am doing,

import tensorflow as tf
from pathlib import Path
import mlflow
import mlflow.keras
from urllib.parse import urlparse
import keras

class Evaluation:
    def __init__(self, config: EvaluationConfig):
        self.config = config

    
    def _valid_generator(self):

        datagenerator_kwargs = dict(
            rescale = 1./255,
            validation_split=0.30
        )

        dataflow_kwargs = dict(
            target_size=self.config.params_image_size[:-1],
            batch_size=self.config.params_batch_size,
            interpolation="bilinear"
        )

        valid_datagenerator = tf.keras.preprocessing.image.ImageDataGenerator(
            **datagenerator_kwargs
        )

        self.valid_generator = valid_datagenerator.flow_from_directory(
            directory=self.config.training_data,
            subset="validation",
            shuffle=False,
            **dataflow_kwargs
        )


    @staticmethod
    def load_model(path: Path) -> tf.keras.Model:
        return tf.keras.models.load_model(path)
    

    def evaluation(self):
        self.model = self.load_model(self.config.path_of_model)
        self._valid_generator()
        self.score = self.model.evaluate(self.valid_generator)
        self.save_score()

    def save_score(self):
        scores = {"loss": self.score[0], "accuracy": self.score[1]}
        save_json(path=Path("scores.json"), data=scores)

    
    def log_into_mlflow(self):
        mlflow.set_registry_uri(self.config.mlflow_uri)
        tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
        
        with mlflow.start_run():
            mlflow.log_params(self.config.all_params)
            mlflow.log_metrics(
                {"loss": self.score[0], "accuracy": self.score[1]}
            )
            
            # Model registry does not work with file store
            if tracking_url_type_store != "file":
            # Register the model with specified flavor options
                flavor_options = {
                    "model_type": keras,
                    "keras_version": tf.keras.__version__,
                    "save_format": "h5"                 
                   }

                # Register the model
                # There are other ways to use the Model Registry, which depends on the use case,
                # please refer to the doc for more information:
                # https://mlflow.org/docs/latest/model-registry.html#api-workflow
                mlflow.keras.log_model(self.model, "model", registered_model_name="VGG16Model", keras_version=tf.keras.__version__,flavor=flavor_options)
            else:
                mlflow.keras.log_model(self.model, "model")

try:
    config = ConfigurationManager()
    eval_config = config.get_evaluation_config()
    evaluation = Evaluation(eval_config)
    evaluation.evaluation()
    evaluation.save_score()
    evaluation.log_into_mlflow()
except Exception as e:
    raise e

Hi @Prajwaal, You can get the keras version using keras.__version__ . Instead of using tf.keras.__version__ could you please try by using keras.__version__. Thank You.