tessl/pypi-torchmetrics

PyTorch native metrics library providing 400+ rigorously tested metrics across classification, regression, audio, image, text, and other ML domains

Overview

Eval results

Files

Information Retrieval Metrics

Name: tessl/pypi-torchmetrics
Author: tessl

Metrics for ranking and retrieval systems including precision at k, mean average precision, and normalized discounted cumulative gain for evaluating search and recommendation systems.

Capabilities

Core Retrieval Metrics

class RetrievalMAP(Metric):
    def __init__(
        self,
        empty_target_action: str = "neg",
        ignore_index: Optional[int] = None,
        **kwargs
    ): ...

class RetrievalMRR(Metric):
    def __init__(
        self,
        empty_target_action: str = "neg",
        ignore_index: Optional[int] = None,
        **kwargs
    ): ...

class RetrievalPrecision(Metric):
    def __init__(
        self,
        empty_target_action: str = "neg",
        ignore_index: Optional[int] = None,
        k: Optional[int] = None,
        **kwargs
    ): ...

class RetrievalRecall(Metric):
    def __init__(
        self,
        empty_target_action: str = "neg",
        ignore_index: Optional[int] = None,
        k: Optional[int] = None,
        **kwargs
    ): ...

Advanced Retrieval Metrics

class RetrievalNormalizedDCG(Metric):
    def __init__(
        self,
        empty_target_action: str = "neg",
        ignore_index: Optional[int] = None,
        k: Optional[int] = None,
        **kwargs
    ): ...

class RetrievalHitRate(Metric):
    def __init__(
        self,
        empty_target_action: str = "neg",
        ignore_index: Optional[int] = None,
        k: Optional[int] = None,
        **kwargs
    ): ...

class RetrievalFallOut(Metric):
    def __init__(
        self,
        empty_target_action: str = "neg",
        ignore_index: Optional[int] = None,
        k: Optional[int] = None,
        **kwargs
    ): ...

Specialized Retrieval Metrics

class RetrievalRPrecision(Metric):
    def __init__(
        self,
        empty_target_action: str = "neg",
        ignore_index: Optional[int] = None,
        **kwargs
    ): ...

class RetrievalAUROC(Metric):
    def __init__(
        self,
        empty_target_action: str = "neg",
        ignore_index: Optional[int] = None,
        **kwargs
    ): ...

class RetrievalPrecisionRecallCurve(Metric):
    def __init__(
        self,
        empty_target_action: str = "neg",
        ignore_index: Optional[int] = None,
        **kwargs
    ): ...

class RetrievalRecallAtFixedPrecision(Metric):
    def __init__(
        self,
        min_precision: float,
        empty_target_action: str = "neg",
        ignore_index: Optional[int] = None,
        **kwargs
    ): ...

Usage Examples

import torch
from torchmetrics.retrieval import (
    RetrievalMAP, RetrievalNormalizedDCG, 
    RetrievalPrecision, RetrievalRecall
)

# Initialize retrieval metrics
map_metric = RetrievalMAP()
ndcg_metric = RetrievalNormalizedDCG(k=10)
precision_k = RetrievalPrecision(k=5)
recall_k = RetrievalRecall(k=10)

# Sample retrieval data
# Scores: higher values indicate higher relevance
preds = torch.tensor([0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1])
# Targets: binary relevance (1 = relevant, 0 = not relevant)
target = torch.tensor([1, 0, 1, 1, 0, 0, 1, 0, 0])
# Indexes: query/user identifiers (same query for all items here)
indexes = torch.zeros(9, dtype=torch.long)

# Compute retrieval metrics
map_score = map_metric(preds, target, indexes=indexes)
ndcg_score = ndcg_metric(preds, target, indexes=indexes)
precision_score = precision_k(preds, target, indexes=indexes)
recall_score = recall_k(preds, target, indexes=indexes)

print(f"MAP: {map_score:.4f}")
print(f"NDCG@10: {ndcg_score:.4f}")
print(f"Precision@5: {precision_score:.4f}")
print(f"Recall@10: {recall_score:.4f}")

Types

RetrievalTarget = Tensor  # Binary relevance labels
RetrievalPreds = Tensor   # Relevance scores
RetrievalIndexes = Tensor # Query/user identifiers
EmptyTargetAction = Union["neg", "pos", "skip"]

Install with Tessl CLI