mbrs.metrics.metricx module

class mbrs.metrics.metricx.MT5ForRegression(config: MT5Config)

Bases: MT5PreTrainedModel

MT5 model for regression.

This implementation is copied from google-research/metricx

config_class

alias of MT5Config

forward(input_ids: LongTensor | None = None, attention_mask: FloatTensor | None = None, decoder_attention_mask: BoolTensor | None = None, head_mask: FloatTensor | None = None, decoder_head_mask: FloatTensor | None = None, cross_attn_head_mask: Tensor | None = None, encoder_outputs: tuple[tuple[Tensor]] | None = None, past_key_values: tuple[tuple[Tensor]] | None = None, inputs_embeds: FloatTensor | None = None, decoder_inputs_embeds: FloatTensor | None = None, labels: FloatTensor | None = None, use_cache: bool | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None) → tuple[Tensor] | MT5ForRegressionOutput

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class mbrs.metrics.metricx.MT5ForRegressionOutput(loss: 'Optional[torch.Tensor]' = None, predictions: 'Optional[torch.Tensor]' = None)

Bases: ModelOutput

loss: Tensor | None = None

predictions: Tensor | None = None

class mbrs.metrics.metricx.MetricMetricX(cfg: Config)

Bases: Metric

MetricX metric class.

References: - MetricX-23: https://aclanthology.org/2023.wmt-1.63 - MetricX-24: https://aclanthology.org/2024.wmt-1.35

Available checkpoints:

• google/metricx-24-hybrid-xxl-v2p6

• google/metricx-24-hybrid-xl-v2p6

• google/metricx-24-hybrid-large-v2p6

• google/metricx-23-xxl-v2p0

• google/metricx-23-xl-v2p0

• google/metricx-23-large-v2p0

• google/metricx-23-qe-xxl-v2p0

• google/metricx-23-qe-xl-v2p0

• google/metricx-23-qe-large-v2p0

class Config(model: str = 'google/metricx-24-hybrid-xxl-v2p6', batch_size: int = 8, fp16: bool = False, bf16: bool = False, cpu: bool = False)

Bases: Config

MetricX 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 = 8

bf16: bool = False

cpu: bool = False

fp16: bool = False

model: str = 'google/metricx-24-hybrid-xxl-v2p6'

HIGHER_IS_BETTER: bool = False

class InputPrefix(hypothesis: 'str', reference: 'str', source: 'str')

Bases: object

hypothesis: str

reference: str

source: str

METRICX23_QE_MODELS = {'google/metricx-23-qe-large-v2p0', 'google/metricx-23-qe-xl-v2p0', 'google/metricx-23-qe-xxl-v2p0'}

METRICX_INPUT_LENGTH_MAP = {MetricXVersion.metricx_23: 1024, MetricXVersion.metricx_24: 1536}

METRICX_INPUT_PREFIX_MAP = {MetricXVersion.metricx_23: MetricMetricX.InputPrefix(hypothesis='candidate: ', reference=' reference: ', source=' source: '), MetricXVersion.metricx_24: MetricMetricX.InputPrefix(hypothesis=' candidate: ', reference=' reference: ', source='source: ')}

METRICX_VERSION_MAP = {'google/metricx-23-large-v2p0': MetricXVersion.metricx_23, 'google/metricx-23-qe-large-v2p0': MetricXVersion.metricx_23, 'google/metricx-23-qe-xl-v2p0': MetricXVersion.metricx_23, 'google/metricx-23-qe-xxl-v2p0': MetricXVersion.metricx_23, 'google/metricx-23-xl-v2p0': MetricXVersion.metricx_23, 'google/metricx-23-xxl-v2p0': MetricXVersion.metricx_23, 'google/metricx-24-hybrid-large-v2p6': MetricXVersion.metricx_24, 'google/metricx-24-hybrid-xl-v2p6': MetricXVersion.metricx_24, 'google/metricx-24-hybrid-xxl-v2p6': MetricXVersion.metricx_24}

class MetricXVersion(value)

Bases: str, Enum

An enumeration.

metricx_23 = 'metricx_23'

metricx_24 = 'metricx_24'

corpus_score(hypotheses: list[str], references_lists: list[list[str]] | None = None, sources: list[str] | None = None) → float

Calculate the corpus-level score.

Parameters:

• hypotheses (list[str]) – Hypotheses.

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

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

Returns:

The corpus score.

Return type:

float

property device: device

Returns the device of the model.

pairwise_scores(hypotheses: list[str], references: list[str], source: str | None = None) → Tensor

Calculate the pairwise scores.

Parameters:

• hypotheses (list[str]) – Hypotheses.

• references (list[str]) – References.

• source (str, optional) – A source.

Returns:

Score matrix of shape (H, R), where H is the number

of hypotheses and R is the number of references.

Return type:

Tensor

score(hypothesis: str, reference: str | None = None, source: str | None = None) → float

Calculate the score of the given hypothesis.

Parameters:

• hypothesis (str) – A hypothesis.

• reference (str, optional) – A reference.

• source (str, optional) – A source.

Returns:

The score of the given hypothesis.

Return type:

float

scorer: MT5ForRegression

scores(hypotheses: list[str], references: list[str] | None = None, sources: list[str] | None = None) → Tensor

Calculate the scores of the given hypothesis.

Parameters:

• hypotheses (list[str]) – N hypotheses.

• references (list[str], optional) – N references.

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

Returns:

The N scores of the given hypotheses.

Return type:

Tensor