File size: 632 Bytes
a6bbf63
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
from transformers import RobertaTokenizerFast, RobertaForSequenceClassification
from utils.commons import HUGGINGFACE_MODEL_PATH, MODEL_SAVE_PATH
import os

class ModelIngestor:
    @staticmethod
    def download_model():
        """Download and save the model"""
        if not os.path.exists(MODEL_SAVE_PATH):
            os.makedirs(MODEL_SAVE_PATH)
            
        RobertaTokenizerFast.from_pretrained(HUGGINGFACE_MODEL_PATH).save_pretrained(MODEL_SAVE_PATH)
        RobertaForSequenceClassification.from_pretrained(HUGGINGFACE_MODEL_PATH).save_pretrained(MODEL_SAVE_PATH)
        print(f"Model saved to {MODEL_SAVE_PATH}")