Spaces:
Running
Running
File size: 1,856 Bytes
4fee431 09c8783 4fee431 |
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 52 53 54 55 |
import os
import shutil
import torch
import torch.nn as nn
import torch.nn.functional as F
import logging
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer, AutoModel
# Configs
REPO_ID = "can-org/Nepali-AI-VS-HUMAN"
BASE_DIR = "./np_text_model"
TOKENIZER_DIR = os.path.join(BASE_DIR, "classifier") # <- update this to match your uploaded folder
WEIGHTS_PATH = os.path.join(BASE_DIR, "model_95_acc.pth") # <- change to match actual uploaded weight
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Define model class
class XLMRClassifier(nn.Module):
def __init__(self):
super(XLMRClassifier, self).__init__()
self.bert = AutoModel.from_pretrained("xlm-roberta-base")
self.classifier = nn.Linear(self.bert.config.hidden_size, 2)
def forward(self, input_ids, attention_mask):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
cls_output = outputs.last_hidden_state[:, 0, :]
return self.classifier(cls_output)
# Globals for caching
_model = None
_tokenizer = None
def download_model_repo():
if os.path.exists(BASE_DIR) and os.path.isdir(BASE_DIR):
logging.info("Model already downloaded.")
return
snapshot_path = snapshot_download(repo_id=REPO_ID)
os.makedirs(BASE_DIR, exist_ok=True)
shutil.copytree(snapshot_path, BASE_DIR, dirs_exist_ok=True)
def load_model():
download_model_repo()
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
model = XLMRClassifier().to(device)
model.load_state_dict(torch.load(WEIGHTS_PATH, map_location=device))
model.eval()
return model, tokenizer
def get_model_tokenizer():
global _model, _tokenizer
if _model is None or _tokenizer is None:
_model, _tokenizer = load_model()
return _model, _tokenizer
|