modules/attention.py

import torch

def get_attention(text, model, tokenizer):
    inputs = tokenizer(text, return_tensors="pt")
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    
    with torch.no_grad():
        outputs = model(**inputs, output_attentions=True)
    
    tokens = [t.replace("Ġ", " ").strip() for t in tokens]
    attentions = [layer[0].numpy() for layer in outputs.attentions]
    
    return tokens, attentions