o %ݫi5@sPdZddlmZmZmZddlZddlmZGdddeZGdddeZ dS) zSProvides a metrics class for the SemDist metric. Authors * Sylvain de Langen 2024 )CallableListLiteralN) MetricStatsc@s^eZdZdZ  ddeeegejfde de fddZ d d Z d d Z dddZddZd S)BaseSemDistStatsuv Base class to implement the SemDist metric, for the variants that estimate a single cosine similarity per pair of target and predicted texts. The SemDist metrics are described by the paper `Evaluating User Perception of Speech Recognition System Quality with Semantic Distance Metric `_. Arguments --------- embed_function : Callable[[List[str]], torch.Tensor] Given a list of sentences, return their summarized embedding using the method of your choice (e.g. mean pooling) scale : float, optional The `α` scale applied to the cosine similarity result for clarity. The default is `1000`, in order to match the authors' recommendation. batch_size : int, optional How many pairs of utterances should be considered at once. Higher is faster but may result in OOM. @@@embed_functionscale batch_sizecCs|||_||_||_dSN)clearr r r )selfr r r rM/home/ubuntu/.local/lib/python3.10/site-packages/speechbrain/utils/semdist.py__init__"s zBaseSemDistStats.__init__cCs"g|_g|_g|_g|_i|_dS)zClears the collected metricsN)ids predictionstargetsscoressummary)rrrrr -s  zBaseSemDistStats.clearcCs(|j||j||j|dS)aN Appends inputs, predictions and targets to internal lists Arguments --------- ids: list the string IDs for the samples predict: list the model's predictions in tokenizable format target: list the ground truths in tokenizable format N)rextendrr)rrpredicttargetrrrappend5s  zBaseSemDistStats.appendNcCsHt |Wdn1swY|dur!|j|S|jS)aSummarize the SemDist metric scores. Performs the actual embedding function call and SemDist calculation. Full set of fields: - `semdist`: The average SemDist over all utterances, multiplied by the scale optionally specified at initialization. Additionally, a `scores` list is populated by this function for each pair of sentences. Each entry of that list is a dict, with the fields: - `key`: the ID of the utterance. - `semdist`: The SemDist of the utterance, multiplied by the scale. Arguments --------- field : str, optional The field to return, if you are only interested in one of them. If specified, a single `float` is returned, otherwise, a dict is. Returns ------- dict from str to float, if `field is None` A dictionary of the fields documented above. float, if `field is not None` The single field selected by `field`. N)torchno_grad_update_summaryr)rfieldrrr summarizeGs   zBaseSemDistStats.summarizec Csd}tdt|j|jD][}|j|||j}|j|||j}|j|||j}||}||}tj j j ||dd}d||j } t |D]\} } |j| | | dqO|| 7}q |t|j} | |jd<dS) zPerforms the actual inference and SemDist estimation, updating the `summary` field. Automatically called by `summarize`.grdimg?)keysemdistr$N)rangelenrr rrr cpurnn functionalcosine_similarityr enumeraterritemsumr) r semdist_sum chunk_idxrref_texthyp_textref_embhyp_emb similarity chunk_semdistiutt_idr$rrrrjs$z BaseSemDistStats._update_summary)rrr )__name__ __module__ __qualname____doc__rrstrrTensorfloatintrr rrrrrrrrs    #rcsFeZdZdZ d dedffdd Zdeedej fd d Z Z S) SemDistStatsaComputes the SemDist metric with a provided HuggingFace Transformers text encoder. Arguments --------- lm : speechbrain.lobes.models.huggingface_transformers.TextEncoder HF Transformers tokenizer and text encoder wrapper to use as a LM. method : "meanpool" or "cls" - `"meanpool"` (default): Computes the mean of all contextualized embeddings, excluding padding tokens. - `"cls"`: Exclusively uses the first contextualized embedding, which with BERT-like tokenizers is the `[CLS]` token, which is typically intended to capture classification information. *args Extra positional arguments passed to the base constructor. **kwargs Extra keyword arguments passed to the base constructor.meanpoolmethod)rAclscs(tj|d|ji|||_||_dS)Nr )superr_embedlmrB)rrFrBargskwargs __class__rrrs zSemDistStats.__init__ sentencesreturncCsdd|D}|j|dd\}}|d}|jdkr7||d}tj|dd}tj|d d|dS|jd krI|d d d d d fStd |jd)aoComputes the LM embedding of a batch of independent sentences, according to the pooling method chosen at initialization. Arguments --------- sentences : list of str List of unprocessed sentences to tokenize and encode. Returns ------- torch.Tensor Embedding of the LM encoder. cSsg|]}d|qS) )join).0sentrrr sz'SemDistStats._embed..T) return_tokensattention_maskrAr r!rCNrzSpecified SemDist method z is invalid)rFr'rB unsqueezerr- ValueError)rrKtokenshiddenmask masked_hiddennonmasked_countsrrrrEs     zSemDistStats._embed)rA) r8r9r:r;rrrr<rr=rE __classcell__rrrIrr@s r@) r;typingrrrrspeechbrain.utils.metric_statsrrr@rrrrs  z