Views
No views yet
max_length for more accurate predictions.1!pip install datasets
2
3import numpy as np
4import torch
5from transformers import BertForSequenceClassification, BertTokenizer
6from datasets import load_dataset
7from sklearn import preprocessing
8
9dataset_id='ahmedheakl/resume-atlas'
10model_id='ahmedheakl/bert-resume-classification'
11label_column = "Category"
12num_labels=43
13output_attentions=False
14output_hidden_states=False
15do_lower_case=True
16add_special_tokens=True
17max_length=512
18pad_to_max_length=True
19return_attention_mask=True
20truncation=True
21
22ds = load_dataset(dataset_id, trust_remote_code=True)
23
24le = preprocessing.LabelEncoder()
25le.fit(ds['train'][label_column])
26
27
28tokenizer = BertTokenizer.from_pretrained(model_id, do_lower_case=do_lower_case)
29model = BertForSequenceClassification.from_pretrained(
30 model_id,
31 num_labels = num_labels,
32 output_attentions = output_attentions,
33 output_hidden_states = output_hidden_states,
34)
35
36model = model.to('cuda').eval()
37sent = ds['train'][0]['Text']
38
39encoded_dict = tokenizer.encode_plus(
40 sent,
41 add_special_tokens=add_special_tokens,
42 max_length=max_length,
43 pad_to_max_length=pad_to_max_length,
44 return_attention_mask=return_attention_mask,
45 return_tensors='pt',
46 truncation=truncation,
47)
48input_ids = encoded_dict['input_ids'].to('cuda')
49attention_mask = encoded_dict['attention_mask'].to('cuda')
50
51outputs = model(
52 input_ids,
53 token_type_ids=None,
54 attention_mask=attention_mask
55)
56
57label_id = np.argmax(outputs['logits'].cpu().detach().tolist(), axis=1)
58print(f'Predicted: {le.inverse_transform(label_id)[0]} | Ground: {ds["train"][0][label_column]}')@article{heakl2024resumeatlas,
title={ResumeAtlas: Revisiting Resume Classification with Large-Scale Datasets and Large Language Models},
author={Heakl, Ahmed and Mohamed, Youssef and Mohamed, Noran and Sharkaway, Ali and Zaky, Ahmed},
journal={arXiv preprint arXiv:2406.18125},
year={2024}
}Heakl, A., Mohamed, Y., Mohamed, N., Sharkaway, A., & Zaky, A. (2024). ResumeAtlas: Revisiting Resume Classification with Large-Scale Datasets and Large Language Models. arXiv (Cornell University). https://doi.org/10.48550/arxiv.2406.18125