components/model_ingestion.py

from transformers import RobertaTokenizerFast, RobertaForSequenceClassification
from utils.commons import HUGGINGFACE_MODEL_PATH, MODEL_SAVE_PATH, setup_logging
import os

logger = setup_logging("model_ingestion.log")

class ModelIngestor:
    @staticmethod
    def download_model():
        """Download and save the model"""
        try:
            if not os.path.exists(MODEL_SAVE_PATH):
                os.makedirs(MODEL_SAVE_PATH)
                logger.info(f"Created directory: {MODEL_SAVE_PATH}")

            logger.info("Downloading tokenizer and model...")
            RobertaTokenizerFast.from_pretrained(HUGGINGFACE_MODEL_PATH).save_pretrained(MODEL_SAVE_PATH)
            RobertaForSequenceClassification.from_pretrained(HUGGINGFACE_MODEL_PATH).save_pretrained(MODEL_SAVE_PATH)
            logger.info(f"Model and tokenizer saved to {MODEL_SAVE_PATH}")
        except Exception as e:
            logger.error(f"Error downloading model: {e}")
            raise