| |
| |
| """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() |
|
|
