Update Main.py

Main.py

@@ -12,16 +12,17 @@ class Main:
         EPOCHS = 6
         NUM_LABELS = 12
         DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        MODEL_PATH = '../input/bert-base-uncased'
-        STATE_DICT = torch.load('../input/ner-resume/model_e6.tar', map_location=DEVICE)
-        tokenizer = BertTokenizerFast('../input/bert-base-uncased/vocab.txt', lowercase=True)
-        MODEL = BertForTokenClassification.from_pretrained(MODEL_PATH, state_dict=STATE_DICT['model_state_dict'], num_labels=12)
+        #MODEL_PATH = '../input/bert-base-uncased'
+        STATE_DICT = torch.load('./models/model_e6.tar', map_location=DEVICE)
+        #tokenizer = BertTokenizerFast('../input/bert-base-uncased/vocab.txt', lowercase=True)
+        tokenizer = BertTokenizerFast.from_pretrained('bert-base-cased')
+        MODEL = BertForTokenClassification.from_pretrained('bert-base-uncased', num_labels=12)
 
         tags_vals = ["UNKNOWN", "O", "Name", "Degree","Skills","College Name","Email Address","Designation","Companies worked at","Graduation Year","Years of Experience","Location"]
         tag2idx = {t: i for i, t in enumerate(tags_vals)}
         idx2tag = {i:t for i, t in enumerate(tags_vals)}
         
-        data = trim_entity_spans(convert_dataturks_to_spacy('../input/resume-entities-for-ner/Entity Recognition in Resumes.json'))
+        data = trim_entity_spans(convert_dataturks_to_spacy('./data/annotated_resumes.json'))
 
         model = MODEL
         model.to(DEVICE);