ss108/legal-citation-bert

This is a NER model meant to be used to detect/extract citations from American legal documents.

Ignore the widget on the model card page; see below for usage.

How to Use the Model

This model outputs token-level predictions, which should be processed as follows to obtain meaningful labels for each token:

from transformers import AutoTokenizer, AutoModelForTokenClassification
import torch

tokenizer = AutoTokenizer.from_pretrained("ss108/legal-citation-bert")
model = AutoModelForTokenClassification.from_pretrained("ss108/legal-citation-bert")

text = "Your example text here"
inputs = tokenizer(text, return_tensors="pt", padding=True)
outputs = model(**inputs)

logits = outputs.logits
predictions = torch.argmax(logits, dim=-1)

tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
predicted_labels = [model.config.id2label[p.item()] for p in predictions[0]]


components = []
for token, label in zip(tokens, predicted_labels):
    components.append(f"{token} : {label}")

concat = " ; ".join(components)
print(concat)