Upload metrics.py with huggingface_hub

metrics.py

@@ -1,7 +1,7 @@
 import re
 import string
 import uuid
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from collections import Counter
 from dataclasses import field
 from typing import Any, Dict, Generator, List, Optional, Tuple
@@ -361,7 +361,7 @@ class BulkInstanceMetric(SingleStreamOperator, MetricWithConfidenceInterval):
  references: List[List[Any]],
  predictions: List[Any],
  additional_inputs: List[Dict],
- ) -> Dict[str, Any]:
+ ) -> List[Dict[str, Any]]:
  pass
 
 
@@ -643,7 +643,6 @@ class HuggingfaceBulkMetric(BulkInstanceMetric):
  predictions: List[str],
  additional_inputs: List[Any],
  ) -> List[Dict[str, Any]]:
- passed_additional_inputs = {}
  passed_additional_inputs = {}
  for additional_input_field in self.hf_additional_input_fields:
  assert (
@@ -1247,7 +1246,7 @@ class SentenceBert(BulkInstanceMetric):
  references: List[List[Any]],
  predictions: List[Any],
  additional_inputs: List[Dict],
- ) -> List[Any]:
+ ) -> List[Dict[str, Any]]:
  scores = []
 
  # we are in a multi-reference case (each prediction may have multiple
@@ -1292,7 +1291,7 @@ class Reward(BulkInstanceMetric):
  references: List[List[Any]],
  predictions: List[Any],
  additional_inputs: List[Dict],
- ) -> List[Any]:
+ ) -> List[Dict[str, Any]]:
  # treat the references as the questions and the predictions as answers
  # assume a single reference
  questions = [refs[0] for refs in references]
@@ -1306,6 +1305,196 @@ class Reward(BulkInstanceMetric):
  return self.pipe(inputs, batch_size=self.batch_size)
 
 
+class Perplexity(BulkInstanceMetric):
+ """Computes the likelihood of generating text Y after text X - P(Y|X)."""
+
+ main_score = "perplexity"
+ reduction_map = {"mean": ["perplexity"]}
+
+ perplexity_prompt: str
+
+ batch_size: int = 32
+ model_name: str
+
+ def compute(
+ self,
+ references: List[List[Any]],
+ predictions: List[Any],
+ additional_inputs: List[Dict],
+ ) -> List[Dict[str, Any]]:
+ """Computes the likelihood of generating text Y after text X - P(Y|X).
+
+ :param references: the list of Y texts as a list of singletons.
+ :param predictions: the list of X texts as a plain list of strings
+
+ :return: the likelihood of generating text Y_i after text X_i = P(Y_i|X_i) for every i.
+ """
+ # make sure all references are singletons
+ assert all(len(ref) == 1 for ref in references)
+
+ # add the instruction as prefix
+ predictions = [f"{self.perplexity_prompt} {x}" for x in predictions]
+ references = [y[0] for y in references]
+
+ # check if the model is enc-dec or dec-only to use the right perplexity computation
+ from transformers import AutoConfig
+
+ config = AutoConfig.from_pretrained(self.model_name, trust_remote_code=True)
+ lm = (
+ self.EncoderDecoderLM(model_name=self.model_name)
+ if config.is_encoder_decoder is True
+ else self.DecoderOnlyLM(model_name=self.model_name)
+ )
+
+ # compute P(Q|P) and store in queue
+ scores = lm.compute_lm(
+ source=predictions, target=references, batch_size=self.batch_size
+ )
+
+ return [{self.main_score: score} for score in scores]
+
+ class AbstractLM(ABC):
+ def __init__(self, model_name):
+ import torch
+ from transformers import AutoTokenizer
+
+ self.model_name = model_name
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+ self.model = self.model_class().from_pretrained(self.model_name)
+ self.is_cuda = torch.cuda.is_available()
+
+ def compute_lm(self, source, target, batch_size: int) -> List[float]:
+ import torch
+
+ scores = []
+
+ with torch.no_grad():
+ # break the documents to batches
+ n_batches = int(len(source) / batch_size)
+ batch_range = range(n_batches + 1)
+ for batch in batch_range:
+ batch_source = source[batch * batch_size : (batch + 1) * batch_size]
+ batch_target = target[batch * batch_size : (batch + 1) * batch_size]
+ if len(batch_source) > 0:
+ # tokenize the source and target
+ tokens_source = self.tokenizer(
+ batch_source, padding=True, return_tensors="pt"
+ )
+ tokens_target = self.tokenizer(
+ batch_target, padding=True, return_tensors="pt"
+ )
+
+ # compute the logits
+ logits, labels = self.compute_batch(
+ tokens_source, tokens_target
+ )
+
+ # the model returns mean over all batch. We run the CE again without reduction
+ # and extarct the mean for each document
+ loss_fct = torch.nn.CrossEntropyLoss(
+ ignore_index=-100, reduction="none"
+ )
+ loss = loss_fct(
+ logits.view(-1, logits.size(-1)), labels.view(-1)
+ )
+ loss = loss.view(len(batch_source), -1)
+
+ # for each document, do mean only over the non zero values (sum(labels>0))
+ batch_loss = torch.sum(loss, dim=1) / torch.sum(
+ labels > 0, dim=1
+ )
+
+ # append the batch scores to the list of all scores
+ scores.append(batch_loss)
+
+ return torch.cat(scores, dim=0).tolist()
+
+ @abstractmethod
+ def model_class(self):
+ pass
+
+ @abstractmethod
+ def compute_batch(self, tokens_source, tokens_target):
+ pass
+
+ class EncoderDecoderLM(AbstractLM):
+ def model_class(self):
+ from transformers import AutoModelForSeq2SeqLM
+
+ return AutoModelForSeq2SeqLM
+
+ def compute_batch(self, tokens_source, tokens_target):
+ tokens_docs_ids = tokens_source["input_ids"]
+ attention = tokens_source["attention_mask"]
+ labels = tokens_target["input_ids"]
+
+ if self.is_cuda:
+ tokens_docs_ids, attention, labels = (
+ tokens_docs_ids.cuda(),
+ attention.cuda(),
+ labels.cuda(),
+ )
+
+ logits = self.model(
+ input_ids=tokens_docs_ids.long(),
+ attention_mask=attention.long(),
+ labels=labels.long(),
+ ).logits
+
+ # replace the padding token in the labels by -100
+ labels[labels == self.tokenizer.pad_token_id] = -100
+
+ return logits, labels
+
+ class DecoderOnlyLM(AbstractLM):
+ def model_class(self):
+ from transformers import AutoModelForCausalLM
+
+ return AutoModelForCausalLM
+
+ def compute_batch(self, tokens_source, tokens_target):
+ import torch
+
+ tokens = torch.cat(
+ [tokens_source["input_ids"], tokens_target["input_ids"]], dim=1
+ )
+ attention = torch.cat(
+ [tokens_source["attention_mask"], tokens_target["attention_mask"]],
+ dim=1,
+ )
+ labels = torch.cat(
+ [
+ torch.zeros_like(tokens_source["input_ids"]).fill_(-100),
+ tokens_target["input_ids"],
+ ],
+ dim=1,
+ )
+
+ # replace the padding token in the labels by -100
+ labels[labels == self.tokenizer.pad_token_id] = -100
+
+ if self.is_cuda:
+ tokens, attention, labels = (
+ tokens.cuda(),
+ attention.cuda(),
+ labels.cuda(),
+ )
+
+ # no need to pass labels as we calculate the loss below per document
+ model_output = self.model(
+ input_ids=tokens.long(), attention_mask=attention.long()
+ )
+ logits = model_output.logits
+
+ # in decoder only, the first token is not being generated, it is taken from the input,
+ # so the model is generating from token 2 to n+1. therefore, we need to skip the last
+ # logit and the first label.
+ shifted_logits = logits[..., :-1, :].contiguous()
+ shifted_labels = labels[..., 1:].contiguous()
+
+ return shifted_logits, shifted_labels
+
+
 class NDCG(GlobalMetric):
  """Normalized Discounted Cumulative Gain: measures the quality of ranking with respect to ground truth ranking scores.