Update Main.py

Main.py

@@ -8,24 +8,24 @@ from transformers import BertTokenizerFast, BertForTokenClassification
 
 class Main:
     def __init__(self):
-        MAX_LEN = 500
-        EPOCHS = 6
-        NUM_LABELS = 12
-        DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.MAX_LEN = 500
+        self.EPOCHS = 6
+        self.NUM_LABELS = 12
+        self.DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         #MODEL_PATH = '../input/bert-base-uncased'
-        STATE_DICT = torch.load('./models/model_e6.tar', map_location=DEVICE)
+        self.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')
+        self.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)}
+        self.tags_vals = ["UNKNOWN", "O", "Name", "Degree","Skills","College Name","Email Address","Designation","Companies worked at","Graduation Year","Years of Experience","Location"]
+        self.tag2idx = {t: i for i, t in enumerate(tags_vals)}
+        self.idx2tag = {i:t for i, t in enumerate(tags_vals)}
         
-        data = self.trim_entity_spans(self.convert_dataturks_to_spacy('./data/annotated_resumes.json'))
+        self.data = self.trim_entity_spans(self.convert_dataturks_to_spacy('./data/annotated_resumes.json'))
 
-        model = MODEL
-        model.to(DEVICE);
+        self.model = MODEL
+        self.model.to(DEVICE);
         
 
     def convert_dataturks_to_spacy(self,dataturks_JSON_FilePath):
@@ -117,7 +117,7 @@ class Main:
         return final_data
 
     def predict(self,test_resume):
-        model.eval()
+        self.model.eval()
         data = self.process_resume(test_resume)
         input_ids, input_mask = self.data['input_ids'].unsqueeze(0), self.data['attention_mask'].unsqueeze(0)
         labels = torch.tensor([1] * input_ids.size(0), dtype=torch.long).unsqueeze(0)