Upload handler.py

handler.py

@@ -0,0 +1,48 @@
+import torch
+from transformers import BertTokenizer, BertForTokenClassification
+
+# Initialize the model and tokenizer
+model_name = "dejanseo/LinkBERT"
+tokenizer = BertTokenizer.from_pretrained(model_name)
+model = BertForTokenClassification.from_pretrained(model_name)
+
+def model_init(path, device='cpu'):
+    """Initialize model."""
+    model.to(device)
+    model.eval()
+    return model
+
+# This function will be called to load the model
+def init():
+    # If your model requires any specific initialization, handle it here
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model_init(model, device=device)
+
+# This function will be called to process requests
+def process(inputs):
+    # Preprocess input data
+    input_data = inputs["inputs"]
+    inputs_tensor = tokenizer(input_data, return_tensors="pt", add_special_tokens=True)
+    input_ids = inputs_tensor["input_ids"]
+
+    # Run model
+    with torch.no_grad():
+        outputs = model(input_ids)
+        predictions = torch.argmax(outputs.logits, dim=-1)
+
+    # Postprocess model outputs
+    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])[1:-1]  # Exclude CLS and SEP tokens
+    predictions = predictions[0][1:-1]
+    result = []
+    for token, pred in zip(tokens, predictions):
+        if pred.item() == 1:
+            result.append(f"<u>{token}</u>")
+        else:
+            result.append(token)
+
+    # Join tokens back into a string
+    reconstructed_text = " ".join(result).replace(" ##", "")
+
+    return {"result": reconstructed_text}
+
+# Note: The actual function signatures for init() and process() might need to be adapted based on Hugging Face's requirements.