Read a torch model from S3

read_torch_from_s3.py

import torch 
from contextlib import contextmanager 
import boto3 
from io import BytesIO 
from transformers import PretrainedConfig, PreTrainedModel 
import json 
from tempfile import NamedTemporaryFile 
 
BUCKET_NAME = "open-jobs-lake" 
 
@contextmanager 
def s3_fileobj(filename): 
    """Yields a file object from the filename at {BUCKET_NAME}/{S3_PATH}"""
    s3 = boto3.client("s3") 
    obj = s3.get_object(Bucket=BUCKET_NAME, Key=filename) 
    yield BytesIO(obj["Body"].read()) 
 
def load_model(s3_path): 
    tempfile = NamedTemporaryFile() 
    with s3_fileobj(f'{s3_path}/pytorch_model.bin') as f: 
        tempfile.write(f.read()) 
 
    with s3_fileobj(f'{s3_path}/config.json') as f: 
        dict_data = json.load(f) 
        config = PretrainedConfig.from_dict(dict_data) 
 
    model = PreTrainedModel.from_pretrained(tempfile.name, config=config) 
    return model 
     
model = load_model('labs/sic/sic4_classifier')