File size: 2,900 Bytes
6f22d12 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | #!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Safe inference example for U4RASD/TypePredictor."""
import argparse
import json
import torch
from modeling_type_predictor import NeoAraBERTTypePredictor
OPEN_MARKER = "[ENT]"
CLOSE_MARKER = "[/ENT]"
def mark_entity(sentence: str, start_char: int, end_char: int) -> str:
entity = sentence[start_char:end_char]
return sentence[:start_char] + " [ENT] " + entity + " [/ENT] " + sentence[end_char:]
def windowed_marked_text(sentence: str, start_char: int, end_char: int, context_chars):
if context_chars is None:
return mark_entity(sentence, start_char, end_char)
left = max(0, start_char - int(context_chars))
right = min(len(sentence), end_char + int(context_chars))
window = sentence[left:right]
return mark_entity(window, start_char - left, end_char - left)
def encode_safely(tokenizer, config, sentence: str, start_char: int, end_char: int):
if not (0 <= start_char < end_char <= len(sentence)):
raise ValueError("Invalid character span.")
open_id = tokenizer.convert_tokens_to_ids(OPEN_MARKER)
close_id = tokenizer.convert_tokens_to_ids(CLOSE_MARKER)
candidates = config.get("context_candidates", [None, 500, 300, 150, 80, 30, 0])
for context in candidates:
text = windowed_marked_text(sentence, start_char, end_char, context)
batch = tokenizer(
text,
return_tensors="pt",
truncation=True,
max_length=int(config.get("max_length", 512)),
)
ids = batch["input_ids"][0].tolist()
if ids.count(open_id) == 1 and ids.count(close_id) == 1 and ids.index(open_id) < ids.index(close_id):
return batch, text, context
raise RuntimeError("Both entity markers could not be preserved after entity-centered truncation.")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", default="U4RASD/TypePredictor")
parser.add_argument("--sentence", required=True)
parser.add_argument("--start-char", type=int, required=True)
parser.add_argument("--end-char", type=int, required=True)
args = parser.parse_args()
model, tokenizer, config = NeoAraBERTTypePredictor.from_pretrained(args.model)
model.eval()
batch, marked_text, context = encode_safely(
tokenizer, config, args.sentence, args.start_char, args.end_char
)
with torch.no_grad():
logits = model(**batch)["logits"]
probabilities = torch.softmax(logits, dim=-1)[0]
index = int(probabilities.argmax())
print(json.dumps({
"entity": args.sentence[args.start_char:args.end_char],
"predicted_type": config["labels"][index],
"confidence": float(probabilities[index]),
"context_chars_used": context,
"marked_text": marked_text,
}, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()
|