mbrs.metrics.comet module

mbrs.metrics.comet module#

class mbrs.metrics.comet.MetricCOMET(cfg: Config)[source]#

Bases: MetricAggregatableCache

COMET metric class.

class Cache(embeddings: Tensor)[source]#

Bases: Cache

Intermediate representations of sentences.

  • embeddings (Tensor): Sentence embeddings of shape (N, D), where N

    is the number of sentences and D is a size of the embedding dimension.

aggregate(reference_lprobs: Tensor | None = None) → Cache[source]#

Aggregate the cached representations.

Parameters:

reference_lprobs (Tensor, optional) – Log-probabilities for each reference sample. The shape must be (len(references),). See https://arxiv.org/abs/2311.05263.

Returns:

An aggregated representation.

Return type:

Cache

cluster(kmeans: Kmeans) → tuple[Cache, Tensor][source]#

Cluster the cached representations.

Parameters:

kmeans (Kmeans) – k-means class to perform clustering.

Returns:

  • Cache: Centroid representations.

  • Tensor: N assigned IDs.

Return type:

tuple[Cache, Tensor]

embeddings: Tensor#
repeat(n: int) → Cache[source]#

Repeat the representations by n times.

Parameters:

n (int) – The number of repetition.

Returns:

The repeated cache.

Return type:

Cache

class Config(model: str = 'Unbabel/wmt22-comet-da', batch_size: int = 64, fp16: bool = False, bf16: bool = False, cpu: bool = False)[source]#

Bases: Config

COMET metric configuration.

  • model (str): Model name or path.

  • batch_size (int): Batch size.

  • fp16 (bool): Use float16 for the forward computation.

  • bf16 (bool): Use bfloat16 for the forward computation.

  • cpu (bool): Use CPU for the forward computation.

batch_size: int = 64#
bf16: bool = False#
cpu: bool = False#
fp16: bool = False#
model: str = 'Unbabel/wmt22-comet-da'#
corpus_score(hypotheses: list[str], references_lists: list[list[str]], sources: list[str] | None = None) → float[source]#

Calculate the corpus-level score.

Parameters:

  • hypotheses (list[str]) – Hypotheses.

  • references_lists (list[list[str]]) – Lists of references.

  • sources (list[str], optional) – Sources.

Returns:

The corpus score.

Return type:

float

Raises:

ValueError – Raise this error when sources are not given.

property device: device#

Returns the device of the model.

property embed_dim: int#

Return the size of embedding dimension.

encode(sentences: list[str]) → Cache[source]#

Encode the given sentences into their intermediate representations.

Parameters:

sentences (list[str]) – Input sentences.

Returns:

Intermediate representations.

Return type:

MetricCOMET.Cache

out_proj(hypotheses_ir: Cache, references_ir: Cache, sources_ir: Cache) → Tensor[source]#

Forward the output projection layer.

Parameters:

  • hypotheses_ir (Cache) – N intermediate representations of hypotheses.

  • references_ir (Cache) – N intermediate representations of references.

  • sources_ir (Cache, optional) – N intermediate representations of sources.

Returns:

N scores.

Return type:

Tensor

Contents