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"}', }