orkungedik/hr-onboaring-doc-classifier

Model Details

Model Description

HR onboarding document type classifier

• Developed by: Orkun Gedik

Training Hyperparameters

• learning_rate=5e-6,
• num_train_epochs=25,
• train_batch_size=16,
• weight_decay=0.01,

Uses

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_v1.1"
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}")

ID2LABEL = { 0: '{"document_type":"adres_belgesi"}', 1: '{"document_type":"banka_belgesi"}', 2: '{"document_type":"askerlik_belgesi"}', 3: '{"document_type":"mezuniyet_belgesi"}', 4: '{"document_type":"nüfus_belgesi"}', 5: '{"document_type":"diğer"}', }