|
from transformers import AutoTokenizer, AutoModelForTokenClassification |
|
import torch |
|
|
|
MODEL_NAME = "ss108/legal-citation-bert" |
|
tokenizer = AutoTokenizer.from_pretrained("ss108/legal-citation-bert") |
|
model = AutoModelForTokenClassification.from_pretrained("ss108/legal-citation-bert") |
|
model.eval() |
|
|
|
def predict(text): |
|
|
|
inputs = tokenizer(text, return_tensors="pt", padding=True) |
|
|
|
|
|
with torch.no_grad(): |
|
outputs = model(**inputs) |
|
logits = outputs.logits |
|
predictions = torch.argmax(logits, dim=-1) |
|
|
|
|
|
labels = [model.config.id2label[pred.item()] for pred in predictions[0]] |
|
|
|
|
|
tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0]) |
|
result = [{'token': token, 'label': label} for token, label in zip(tokens, labels) if token not in tokenizer.all_special_tokens] |
|
|
|
return result |