Update app.py

app.py

@@ -12,23 +12,27 @@ from pymongo import MongoClient
 from transformers import AutoTokenizer
 tokenizer = AutoTokenizer.from_pretrained("law-ai/InLegalBERT")
 
-model_preamble = torch.load("nerbert_preamble.pt")
-model_judgment = torch.load("nerbert.pt")
-
 from transformers import BertForTokenClassification
-# class BertModel(torch.nn.Module):
+class BertModel(torch.nn.Module):
+
+    def __init__(self):
 
-#     def __init__(self):
+        super(BertModel, self).__init__()
 
-#         super(BertModel, self).__init__()
+        self.bert = BertForTokenClassification.from_pretrained('law-ai/InLegalBERT', num_labels=len(unique_labels))
 
-#         self.bert = BertForTokenClassification.from_pretrained('law-ai/InLegalBERT', num_labels=len(unique_labels))
+    def forward(self, input_id, mask, label):
 
-#     def forward(self, input_id, mask, label):
+        output = self.bert(input_ids=input_id, attention_mask=mask, labels=label, return_dict=False)
 
-#         output = self.bert(input_ids=input_id, attention_mask=mask, labels=label, return_dict=False)
+        return output
+
+model_preamble = BertModel()
+model_preamble = torch.load("nerbert_preamble.pt")
+
+model_judgment = BertModel()
+model_judgment = torch.load("nerbert.pt")
 
-#         return output
 
 unique_labels_preamble = {'I-PETITIONER', 'I-COURT', 'B-COURT', 'B-JUDGE', 'I-LAWYER', 'B-RESPONDENT', 'I-JUDGE', 'B-PETITIONER', 'I-RESPONDENT', 'B-LAWYER', 'O'}
 unique_labels_judgment = {'B-WITNESS', 'I-PETITIONER', 'I-JUDGE', 'B-STATUTE', 'B-OTHER_PERSON', 'B-CASE_NUMBER', 'I-ORG', 'I-PRECEDENT', 'I-RESPONDENT', 'B-PROVISION', 'O', 'I-WITNESS', 'B-ORG', 'I-COURT', 'B-RESPONDENT', 'I-DATE', 'B-GPE', 'I-CASE_NUMBER', 'B-DATE', 'B-PRECEDENT', 'I-GPE', 'B-COURT', 'B-JUDGE', 'I-STATUTE', 'B-PETITIONER', 'I-OTHER_PERSON', 'I-PROVISION'}