from datasets import load_dataset
from transformers import ViTImageProcessor, ViTForImageClassification
import torch

# Convert the image to RGB
example = example["image"].convert('RGB')

model_name = "orkungedik/hr-onboaring-doc-classifier"
processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)

inputs = processor(images=example, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits

predicted_class_idx = logits.argmax(-1).item()
label = model.config.id2label[predicted_class_idx]

print(f"Predicted class: {label}")

probs = torch.nn.functional.softmax(logits, dim=-1)
top5 = torch.topk(probs, 5)

for i in range(5):
    idx = top5.indices[0][i].item()
    prob = top5.values[0][i].item()
    print(f"{model.config.id2label[idx]}: {prob:.4f}")

orkungedik/hr-onboaring-doc-classifier_v1.0

Model Details

Model Description

Training Hyperparameters

Uses

orkungedik/
hr-onboaring-doc-classifier_v1.0