Update custom_model.py

custom_model.py

@@ -1,19 +1,21 @@
-# custom_model.py
-from transformers import BertTokenizerFast, BertForTokenClassification
+from transformers import BertForTokenClassification, BertTokenizerFast
 from peft import PeftModel
 import torch
 
 class CustomBertForTokenClassification(BertForTokenClassification):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
+    def __init__(self, config):
+        super().__init__(config)
+        
     def predict_entities(self, sentence, label_mapping):
         # Tokenize the input sentence
-        tokenizer = BertTokenizerFast.from_pretrained(self.config._name_or_path)
+        tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
         inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding=True)
-        inputs = {key: value.to(self.device) for key, value in inputs.items()}
         
-        # Get predictions
+        # Move inputs to device
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        inputs = {key: value.to(device) for key, value in inputs.items()}
+        
+        # Get model predictions
         with torch.no_grad():
             outputs = self(**inputs)