pip install "alquimia-fair-forge[explainability]"

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Disable autocast for attribution compatibility
torch.set_autocast_enabled(False)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    torch_dtype=torch.float16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")

messages = [
    {"role": "system", "content": "Answer concisely."},
    {"role": "user", "content": "What is photosynthesis?"}
]

# Use the tokenizer's chat template
prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

from fair_forge.explainability import AttributionExplainer, Lime, Granularity

explainer = AttributionExplainer(
    model=model,
    tokenizer=tokenizer,
    default_method=Lime,
    default_granularity=Granularity.WORD,
    verbose=True,
)

result = explainer.explain(
    prompt=prompt,
    target="Photosynthesis is the process plants use to convert sunlight into energy.",
)

# Get top contributing words
print("Top 10 Most Important Words:")
for attr in result.get_top_k(10):
    print(f"  '{attr.text}': {attr.score:+.4f}")

# Access all attributions
for attr in result.attributions:
    print(f"Position {attr.position}: '{attr.text}' = {attr.score:.4f}")

from fair_forge.explainability import (
    Lime, Occlusion, Saliency, IntegratedGradients
)

# LIME (perturbation-based, recommended default)
result_lime = explainer.explain(prompt=prompt, target=target, method=Lime)

# Occlusion (perturbation-based)
result_occ = explainer.explain(prompt=prompt, target=target, method=Occlusion)

# Saliency (gradient-based, faster)
result_sal = explainer.explain(prompt=prompt, target=target, method=Saliency)

# Integrated Gradients (gradient-based, more accurate)
result_ig = explainer.explain(prompt=prompt, target=target, method=IntegratedGradients)

# Prepare batch (prompts must be pre-formatted)
items = [
    (
        tokenizer.apply_chat_template(
            [{"role": "user", "content": "What is AI?"}],
            tokenize=False, add_generation_prompt=True
        ),
        "AI is artificial intelligence."
    ),
    (
        tokenizer.apply_chat_template(
            [{"role": "user", "content": "What is ML?"}],
            tokenize=False, add_generation_prompt=True
        ),
        "ML is machine learning, a subset of AI."
    ),
]

# Process batch
batch_results = explainer.explain_batch(items)

print(f"Processed {len(batch_results)} items")
print(f"Total time: {batch_results.total_compute_time_seconds:.2f}s")

for i, result in enumerate(batch_results):
    print(f"\nItem {i+1}: Top words = {[a.text for a in result.get_top_k(3)]}")

class AttributionResult(BaseModel):
    prompt: str                      # Input prompt
    target: str                      # Target output
    method: AttributionMethod        # Method used
    granularity: Granularity         # Granularity level
    attributions: list[TokenAttribution]  # Token scores
    metadata: dict[str, Any]         # Compute time, etc.

class TokenAttribution(BaseModel):
    text: str              # Token/word/sentence text
    score: float           # Attribution score (can be negative)
    position: int          # Position in sequence
    normalized_score: float | None  # Score normalized to [0, 1]

# Get top K most important tokens
top_tokens = result.get_top_k(10)

# Get all attributions sorted by importance
sorted_attrs = result.top_attributions

# Export for visualization
viz_data = result.to_dict_for_visualization()
# Returns: {"tokens": [...], "scores": [...], "normalized_scores": [...]}

# Export as dict (JSON-serializable)
result_dict = result.model_dump()

# Display interactive visualization
explainer.visualize(result)

# Get HTML string
html = explainer.visualize(result, return_html=True)

# Use in web apps, save to file, etc.
with open("attribution.html", "w") as f:
    f.write(html)

from fair_forge.explainability import (
    AttributionExplainer,
    AttributionResultParser,
)

class MyCustomParser(AttributionResultParser):
    def parse(self, raw_result) -> tuple[list[str], list[float]]:
        # Extract tokens and scores from your custom format
        tokens = raw_result["words"]
        scores = raw_result["importance"]
        return tokens, scores

# Use custom parser
explainer = AttributionExplainer(
    model=model,
    tokenizer=tokenizer,
    result_parser=MyCustomParser(),
)

from fair_forge.explainability import compute_attributions, Lime

result = compute_attributions(
    model=model,
    tokenizer=tokenizer,
    prompt=prompt,
    target=target,
    method=Lime,
    granularity=Granularity.WORD,
)

print("Analysis of attribution results:")
print("=" * 50)

# Positive contributors
positive = [a for a in result.attributions if a.score > 0]
print(f"\nPositive contributors ({len(positive)} tokens):")
for attr in sorted(positive, key=lambda x: x.score, reverse=True)[:5]:
    print(f"  '{attr.text}': +{attr.score:.4f}")

# Negative contributors
negative = [a for a in result.attributions if a.score < 0]
print(f"\nNegative contributors ({len(negative)} tokens):")
for attr in sorted(negative, key=lambda x: x.score)[:5]:
    print(f"  '{attr.text}': {attr.score:.4f}")