codeby-hp/finetuneTinyBERT-SentimentClassification

1import boto3
2import logging
3
4import os
5from dotenv import load_dotenv
6
7load_dotenv()
8
9logging.basicConfig(level=logging.INFO)
10logger = logging.getLogger(__name__)
11
12
13def download_model_from_s3(
14    local_dir="./model", s3_prefix="ml-models/tinybert-sentiment-analysis"
15):
16    """
17    Download the fine-tuned model from S3 bucket
18    """
19    bucket_name = os.getenv("BUCKET_NAME")
20    if not bucket_name:
21        raise ValueError("BUCKET_NAME not found in .env file")
22
23    os.makedirs(local_dir, exist_ok=True)
24
25    s3_client = boto3.client("s3")
26
27    model_files = [
28        "config.json",
29        "model.safetensors",
30        "special_tokens_map.json",
31        "tokenizer_config.json",
32        "tokenizer.json",
33        "vocab.txt",
34    ]
35
36    logger.info(f"Downloading model from S3 bucket: {bucket_name}/{s3_prefix}")
37
38    for file_name in model_files:
39        try:
40            local_file_path = os.path.join(local_dir, file_name)
41
42            if os.path.exists(local_file_path):
43                logger.info(f"File {file_name} already exists, skipping...")
44                continue
45
46            s3_key = f"{s3_prefix}/{file_name}" if s3_prefix else file_name
47
48            logger.info(f"Downloading {s3_key}...")
49            s3_client.download_file(bucket_name, s3_key, local_file_path)
50            logger.info(f"Successfully downloaded {file_name}")
51        except Exception as e:
52            logger.error(f"Error downloading {file_name}: {e}")
53            raise
54
55    logger.info("Model download completed successfully")
56    return local_dir
57