src/save_model.py

import os

def save_model_and_tokenizer(model, tokenizer, save_directory):
    """

    Save model and tokenizer to the specified directory.



    Args:

    - model: The model to save.

    - tokenizer: The tokenizer to save.

    - save_directory: Directory where the model and tokenizer should be saved.

    """
    try:
        # Ensure the save directory exists
        os.makedirs(save_directory, exist_ok=True)

        # Save model and tokenizer
        model.save_pretrained(save_directory, safe_serialization=True)
        tokenizer.save_pretrained(save_directory)

        print(f"Model and tokenizer saved locally at {save_directory}")
    except Exception as e:
        print(f"Error saving model and tokenizer: {str(e)}")
        raise