from transformers import AutoModelForCausalLM, AutoTokenizer
from fair_forge.explainability import AttributionExplainer, Lime, Granularity

# Load model
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")

# Format prompt (user responsibility - varies by model)
messages = [{"role": "user", "content": "What is gravity?"}]
prompt = tokenizer.apply_chat_template(messages, tokenize=False)

# Compute attributions
explainer = AttributionExplainer(model, tokenizer)
result = explainer.explain(
    prompt=prompt,
    target="Gravity is the force of attraction between objects.",
    method=Lime,
)

# Show most important words
for attr in result.get_top_k(5):
    print(f"'{attr.text}': {attr.score:.4f}")