Spaces:
Sleeping
Sleeping
File size: 2,012 Bytes
022acf4 |
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
from typing import Tuple
import numpy as np
from sklearn.metrics import (accuracy_score, f1_score, precision_score,
recall_score)
from tqdm.auto import tqdm
import torch
from newsclassifier.config.config import Cfg, logger
from newsclassifier.data import NewsDataset, collate
from newsclassifier.models import CustomModel
from torch.utils.data import DataLoader
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def test_step(test_loader: DataLoader, model) -> Tuple[np.ndarray, np.ndarray]:
"""Eval step."""
model.eval()
y_trues, y_preds = [], []
with torch.inference_mode():
for step, (inputs, labels) in tqdm(enumerate(test_loader)):
inputs = collate(inputs)
for k, v in inputs.items():
inputs[k] = v.to(device)
labels = labels.to(device)
y_pred = model(inputs)
y_trues.extend(labels.cpu().numpy())
y_preds.extend(torch.argmax(y_pred, dim=1).cpu().numpy())
return np.vstack(y_trues), np.vstack(y_preds)
def inference():
logger.info("Loading inference data.")
try:
test_dataset = NewsDataset(os.path.join(Cfg.preprocessed_data_path, "test.csv"))
test_loader = DataLoader(test_dataset, batch_size=Cfg.batch_size, shuffle=False, num_workers=4, pin_memory=True, drop_last=False)
except Exception as e:
logger.error(e)
logger.info("loading model.")
try:
model = CustomModel(num_classes=Cfg.num_classes)
model.load_state_dict(torch.load(Cfg.model_path, map_location=torch.device("cpu")))
model.to(device)
except Exception as e:
logger.error(e)
y_true, y_pred = test_step(test_loader, model)
print(
f'Precision: {precision_score(y_true, y_pred, average="weighted")} \n Recall: {recall_score(y_true, y_pred, average="weighted")} \n F1: {f1_score(y_true, y_pred, average="weighted")} \n Accuracy: {accuracy_score(y_true, y_pred)}'
)
|