Spaces:
Running
Running
File size: 1,517 Bytes
6f034a7 09c8783 6f034a7 4fee431 6f034a7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import os
import shutil
import logging
from transformers import GPT2LMHeadModel, GPT2TokenizerFast, GPT2Config
from huggingface_hub import snapshot_download
import torch
from dotenv import load_dotenv
load_dotenv()
REPO_ID = "can-org/AI-Content-Checker"
MODEL_DIR = "./models"
TOKENIZER_DIR = os.path.join(MODEL_DIR, "model")
WEIGHTS_PATH = os.path.join(MODEL_DIR, "model_weights.pth")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
_model, _tokenizer = None, None
def warmup():
global _model, _tokenizer
# Ensure punkt is available
download_model_repo()
_model, _tokenizer = load_model()
logging.info("Its ready")
def download_model_repo():
if os.path.exists(MODEL_DIR) and os.path.isdir(MODEL_DIR):
logging.info("Model already exists, skipping download.")
return
snapshot_path = snapshot_download(repo_id=REPO_ID)
os.makedirs(MODEL_DIR, exist_ok=True)
shutil.copytree(snapshot_path, MODEL_DIR, dirs_exist_ok=True)
def load_model():
tokenizer = GPT2TokenizerFast.from_pretrained(TOKENIZER_DIR)
config = GPT2Config.from_pretrained(TOKENIZER_DIR)
model = GPT2LMHeadModel(config)
model.load_state_dict(torch.load(WEIGHTS_PATH, map_location=device))
model.to(device)
model.eval()
return model, tokenizer
def get_model_tokenizer():
global _model, _tokenizer
if _model is None or _tokenizer is None:
download_model_repo()
_model, _tokenizer = load_model()
return _model, _tokenizer
|