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):
    # Tokenize the input text
    inputs = tokenizer(text, return_tensors="pt", padding=True)

    # Get model predictions
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predictions = torch.argmax(logits, dim=-1)

    # Convert predictions to labels
    labels = [model.config.id2label[pred.item()] for pred in predictions[0]]

    # Align labels with tokens
    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