Update handler.py

handler.py

@@ -1,14 +1,23 @@
-from transformers import AutoModelForTokenClassification, AutoTokenizer
+from transformers import BertForTokenClassification, BertTokenizer, AutoConfig
 import torch
 from typing import Dict, List, Any
 
 class EndpointHandler:
     def __init__(self, path: str = "dejanseo/LinkBERT"):
-        # Initialize tokenizer and model with the specified path
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModelForTokenClassification.from_pretrained(path)
+        # Load the configuration from the saved model
+        self.config = AutoConfig.from_pretrained(path)
+
+        # Make sure to specify the correct model name for bert-large-cased
+        # Adjust num_labels according to your model's configuration
+        self.model = BertForTokenClassification.from_pretrained(
+            path,
+            config=self.config
+        )
         self.model.eval()  # Set model to evaluation mode
 
+        # Load the tokenizer for bert-large-cased
+        self.tokenizer = BertTokenizer.from_pretrained("bert-large-cased")
+
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         # Extract input text from the request
         inputs = data.get("inputs", "")
@@ -29,7 +38,7 @@ class EndpointHandler:
         # Reconstruct the text with annotations for token classification
         result = []
         for token, pred in zip(tokens, predictions):
-            if pred == 1:  # Assuming '1' is the label for the class of interest
+            if pred == 1:  # Adjust this based on your classification needs
                 result.append(f"<u>{token}</u>")
             else:
                 result.append(token)
@@ -39,4 +48,5 @@ class EndpointHandler:
         # Return the processed text in a structured format
         return [{"text": reconstructed_text}]
 
-# Note: You'll need to replace 'path' with the actual path or identifier of your model when initializing the EndpointHandler.
+# Note: Ensure the path "dejanseo/LinkBERT" is correctly pointing to your model's location
+# If the model is locally saved, adjust the path accordingly