Add files using upload-large-folder tool

6b56858 verified 6 days ago

16.8 kB

import numpy as np
from math import log
from collections import defaultdict
import math
import os
import typing

import torch
import torch.nn.functional as F
import torchmetrics
import transformers
from transformers import logging
logging.set_verbosity_error()

# TODO: for tc

LOG2 = math.log(2)

# For take the total correlation between x0, x1, and y


class TC():
    def __init__(self):
        self.x0 = []
        self.x1 = torch.tensor([])

    def update(self, x0, x1):
        # note that x0, x1, y are all torch tensors
        assert len(x0) == len(x1), "All must have same length"

        # it will be a [8, 1024]
        if self.x1.device != x1.device:
            self.x1 = self.x1.to(x1.device)
        self.x0 += x0
        self.x1 = torch.cat((self.x1, x1), dim=0)

    def entropy_from_counts(self, counts, base=2):
        """
        Compute the empirical entropy H(X) from a dictionary of counts,
        using either base-2 (bits) or base-e (nats).
        """
        total = sum(counts.values())
        if total == 0:
            return 0.0
        entropy = 0.0
        for c in counts.values():
            if c > 0:
                p = c / total
                # Use math.log(p, base) for bits (base=2) or nats (base=np.e)
                entropy -= p * log(p, base)
        return entropy

    def compute_total_correlation_x1(self, x1_group):
        """
        Given an array x1_group of shape (N, 16, 16), where each row is
        one sample of x1, estimate T(X1) = sum_j H(X1_j) - H(X1)
        by naive frequency counting.
        We treat each of the 256 pixels in the 16x16 image as a separate
        discrete variable. We compute:
           H(X1)   via joint frequencies over all 256 pixels
           H(X1_j) via marginal frequencies of each pixel j
        Then T(X1) = sum_j H(X1_j) - H(X1).
        """
        N = x1_group.shape[0]
        if N < 2:
            # With <2 samples, you cannot really estimate correlation reliably.
            # We return 0.0 by convention or skip it entirely.
            return 0.0

        # Flatten each image into 256-dimensional vector:
        # shape becomes (N, 1024)
        flattened = x1_group.reshape(N, -1)
        d = flattened.shape[1]  # should be 1024

        # ---- Joint distribution (all 256 dims) ----#
        # We'll store each flattened row as a tuple, then count frequencies.
        joint_counts = defaultdict(int)
        for row in flattened:
            key = tuple(row)   # row is length-1024
            joint_counts[key] += 1
        H_joint = self.entropy_from_counts(joint_counts, base=2)

        # ---- Marginal distributions (one pixel at a time) ----#
        # We'll compute an entropy for each pixel dimension j.
        marginal_entropies = []
        for j in range(d):
            counts_j = defaultdict(int)
            for row in flattened:
                val_j = row[j].long().item()
                counts_j[val_j] += 1
            H_j = self.entropy_from_counts(counts_j, base=2)
            marginal_entropies.append(H_j)

        sum_marginals = sum(marginal_entropies)

        # ---- Total correlation ----#
        T_val = sum_marginals - H_joint
        return T_val, sum_marginals, H_joint

    def compute_conditional_total_correlation_x1_given_x0y(self, x0, x1):
        """
        1) Group all samples by unique x0[i] pair.
           - Here x0[i] is a [1024] tensor
        2) Collect the corresponding x1[i] arrays for each group.
        3) For each group, estimate T(X1) by naive frequency counting
           (thus approximating T(X1 | x0=x0_val, y=y_val)).
        Returns a dict:  (x0_bytes) -> estimated total correlation in bits.
        """

        # Safety checks:
        assert len(x0) == len(x1), "All must have same length"
        N = len(x0)

        # Group x1 by unique x0
        groups = defaultdict(list)
        for i in range(N):
            # We need a hashable key for x0[i], which is shape (1028)
            # Convert to bytes, or a tuple if you prefer
            x0_key = x0[i]
            key = x0_key
            groups[key].append(x1[i])

        # Compute total correlation for each group
        results = {}
        marginals = {}
        joints = {}
        for key, x0_x1_tup in groups.items():
            # x1_list is a tensor, shape each (1024)
            # print(len(x1_list))
            # shape = (num_samples_for_this_group, 1024)
            group = torch.stack(x0_x1_tup, dim=0)
            T_val, sum_marginals, H_joint = self.compute_total_correlation_x1(
                group)
            results[key] = T_val
            marginals[key] = sum_marginals
            joints[key] = H_joint

        return results, joints, marginals

    def compute(self):
        """
        Compute the total correlation of x1 given x0 and y.
        """
        tc_results, joints, marginals = self.compute_conditional_total_correlation_x1_given_x0y(
            self.x0, self.x1)

        avg_tc = np.mean(list(tc_results.values()))
        avg_joints = np.mean(list(joints.values()))
        avg_marginals = np.mean(list(marginals.values()))

        return avg_tc, avg_joints, avg_marginals


class NLL(torchmetrics.aggregation.MeanMetric):
    def update(self,
               value: typing.Union[float, torch.Tensor],
               weight: typing.Union[float, torch.Tensor] = 1.0) -> None:
        """Update state with data.
        Args:
          value: Either a float or tensor containing data.
            Additional tensor dimensions will be flattened
          weight: Either a float or tensor containing weights
            for calculating the average. Shape of weight should
            be able to broadcast with the shape of `value`.
            Default to `1.0` corresponding to simple harmonic
            average.
        """
        # broadcast weight to value shape
        if not isinstance(value, torch.Tensor):
            value = torch.as_tensor(value, dtype=self.dtype,
                                    device=self.device)
        if (weight is not None
                and not isinstance(weight, torch.Tensor)):
            weight = torch.as_tensor(weight,
                                     dtype=self.dtype,
                                     device=self.device)
        weight = torch.broadcast_to(weight, value.shape)
        value, weight = self._cast_and_nan_check_input(value,
                                                       weight)

        if value.numel() == 0:
            return
        self.mean_value += value.sum()
        self.weight += weight.sum()


class BPD(NLL):
    def compute(self) -> torch.Tensor:
        """Computes the bits per dimension.
        Returns:
          bpd
        """
        return self.mean_value / self.weight / LOG2


class Perplexity(NLL):
    def compute(self) -> torch.Tensor:
        """Computes the Perplexity.
        Returns:
         Perplexity
        """
        return torch.exp(self.mean_value / self.weight)


class Metrics:
    def __init__(self, gen_ppl_eval_model_name_or_path=None,
                 eval_ppl_batch_size=None) -> None:
        metrics = torchmetrics.MetricCollection({
            'nll': NLL(), 'bpd': BPD(), 'ppl': Perplexity()})
        metrics.set_dtype(torch.float64)
        self.train_nlls = metrics.clone(prefix='train/')
        self.train_aux = BPD()
        self.valid_nlls = metrics.clone(prefix='val/')
        self.valid_aux = BPD()
        self.gen_ppl = Perplexity()
        self.sample_entropy = torchmetrics.aggregation.MeanMetric()
        self.unique_token_count = 0
        self.tc = TC()
        self.eval_ppl_batch_size = eval_ppl_batch_size
        self.gen_ppl_eval_model_name_or_path = gen_ppl_eval_model_name_or_path
        self.tokenizer = transformers.AutoTokenizer.\
            from_pretrained(gen_ppl_eval_model_name_or_path)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id

    def to(self, *args, **kwargs):
        self.gen_ppl = self.gen_ppl.to(*args, **kwargs)
        self.sample_entropy = self.sample_entropy.to(*args, **kwargs)
        self.train_nlls = self.train_nlls.to(*args, **kwargs)
        self.train_aux = self.train_aux.to(*args, **kwargs)
        self.valid_nlls = self.valid_nlls.to(*args, **kwargs)
        self.valid_aux = self.valid_aux.to(*args, **kwargs)

    def reset(self):
        self.gen_ppl.reset()
        self.sample_entropy.reset()
        self.train_nlls.reset()
        self.train_aux.reset()
        self.valid_nlls.reset()
        self.valid_aux.reset()

    def update_train(self, nll, aux_loss, num_tokens):
        self.train_nlls.update(nll, num_tokens)
        self.train_aux.update(aux_loss, num_tokens)

    def update_valid(self, nll, aux_loss, num_tokens):
        self.valid_nlls.update(nll, num_tokens)
        self.valid_aux.update(aux_loss, num_tokens)

    @torch.no_grad()
    def _eval_retokenize(self, text_samples, max_length,
                         device):
        """Retokenizes samples for the eval model.
        Args:
            text_samples: List of sentences generated by the model.
        Returns:
            samples: Samples re-tokenized for the eval model
            attn_mask: Attention mask for the eval model
            eval_context_size: Size of the context for the eval model
        """
        if 'llama2' in self.gen_ppl_eval_model_name_or_path.lower():
            tokenizer_kwargs = {
                'text_samples': text_samples,
                'return_tensors': 'pt',
                'return_token_type_ids': False,
                'return_attention_mask': True,
                'truncation': True,
                'padding': True,
                'max_length': max_length,
            }
            eval_context_size = 4096
        elif 'llama3' in self.gen_ppl_eval_model_name_or_path.lower():
            tokenizer_kwargs = {
                'text': text_samples,
                'return_tensors': 'pt',
                'padding': True,
            }
            eval_context_size = 8192
        else:
            tokenizer_kwargs = {
                'return_tensors': 'pt',
                'return_token_type_ids': False,
                'return_attention_mask': True,
                'truncation': True,
                'padding': True,
                'max_length': max_length,
            }
            eval_context_size = 1024
        samples = self.tokenizer(text_samples,
                                 **tokenizer_kwargs)
        attn_mask = samples['attention_mask']
        samples = samples['input_ids']
        if 'llama' not in self.gen_ppl_eval_model_name_or_path.lower():
            attn_mask = attn_mask.to(device)
            samples = samples.to(device)
        return samples, attn_mask, eval_context_size

    @torch.no_grad()
    def record_entropy(self, tokens):
        for sample in tokens:
            _, counts = torch.unique(
                sample, return_counts=True, sorted=False)
            entropy = torch.special.entr(
                counts.float() / counts.sum()).sum().item()
            self.sample_entropy.update(entropy)

    @torch.no_grad()
    def record_unique_tokens(self, samples):
        """Record the count of unique tokens across all samples.
        
        Args:
            samples: torch.Tensor of shape (batch_size, seq_length)
        
        Returns:
            int: Number of unique tokens in the samples
        """
        unique_tokens = torch.unique(samples.flatten())
        self.unique_token_count = len(unique_tokens)
        return self.unique_token_count

    def reset_unique_tokens(self):
        """Reset unique token count."""
        self.unique_token_count = 0

    @torch.no_grad()
    def record_generative_perplexity(
            self,
            text_samples: typing.List[str],
            max_length: int,
            retokenize: bool = True,
            device='cuda') -> None:

        os.environ['TOKENIZERS_PARALLELISM'] = 'false'
        if 'llama' not in self.gen_ppl_eval_model_name_or_path:
            eval_model = transformers.AutoModelForCausalLM.from_pretrained(
                self.gen_ppl_eval_model_name_or_path).eval()
            eval_model = eval_model.to(device)
            # Re-tokenize using eval model's tokenizer
            if retokenize:
                (samples, attn_mask,
                 eval_context_size) = self._eval_retokenize(
                    text_samples, max_length=max_length, device=device)
            else:
                samples = text_samples
                attn_mask = torch.ones(samples.shape).to(device)
                eval_context_size = samples.shape[-1]
            batch_size = min(self.eval_ppl_batch_size,
                             samples.shape[0])
            num_batches = samples.shape[0] // batch_size
            for i in range(num_batches):
                _samples = torch.split(
                    samples[i * batch_size: (i + 1) * batch_size],
                    eval_context_size,
                    dim=-1)
                _attn_mask = torch.split(
                    attn_mask[i * batch_size: (i + 1) * batch_size],
                    eval_context_size,
                    dim=-1)
                for (sample_chunk, attn_mask_chunk) in zip(_samples,
                                                           _attn_mask):
                    logits = eval_model(sample_chunk.to(device),
                                        attention_mask=attn_mask_chunk.to(device))
                    logits = logits[0].transpose(-1, -2)
                    nlls = F.cross_entropy(logits[..., :-1],
                                           sample_chunk[..., 1:],
                                           reduction='none')
                    first_eos = (
                        sample_chunk
                        == self.tokenizer.eos_token_id).cumsum(-1) == 1
                    token_mask = sample_chunk != self.tokenizer.eos_token_id
                    valid_tokens = first_eos[..., 1:] + token_mask[..., 1:]
                    self.gen_ppl.update(nlls * valid_tokens, valid_tokens)
        else:
            eval_model = transformers.AutoModelForCausalLM.from_pretrained(
                self.gen_ppl_eval_model_name_or_path,
                torch_dtype=torch.bfloat16).eval()
            eval_model = eval_model.to(device)
            # Re-tokenize using eval model's tokenizer
            tokenizer_llama = transformers.AutoTokenizer.from_pretrained(
                self.gen_ppl_eval_model_name_or_path)
            if tokenizer_llama.pad_token is None:
                tokenizer_llama.pad_token = tokenizer_llama.eos_token
                tokenizer_llama.pad_token_id = tokenizer_llama.eos_token_id
            tokenizer_gpt = transformers.AutoTokenizer.from_pretrained(
                'gpt2')
            if tokenizer_gpt.pad_token is None:
                tokenizer_gpt.pad_token = tokenizer_gpt.eos_token
                tokenizer_gpt.pad_token_id = tokenizer_gpt.eos_token_id
            num_samples = len(text_samples)
            batch_size = min(16, num_samples)
            num_batches = num_samples // batch_size

            with torch.inference_mode():
                # 1. divide into batches of 16
                # 2. encode each batch
                # 3. eval each batch of 16
                for i in range(num_batches):
                    batch_text_samples = text_samples[i *
                                                      batch_size:(i+1)*batch_size]
                    encoded_inputs = tokenizer_llama(
                        batch_text_samples,
                        return_tensors="pt",
                        padding=True,
                    )

                    input_ids = encoded_inputs['input_ids'].to(device)
                    attention_mask = encoded_inputs['attention_mask'].to(
                        device)

                    labels = input_ids.clone()
                    labels[labels == tokenizer_llama.pad_token_id] = 50000
                    labels = labels.to(device)

                    outputs = eval_model(
                        input_ids=input_ids, attention_mask=attention_mask, labels=labels)

                    llama_logits = outputs.logits

                    logits = llama_logits.transpose(-1, -2)
                    nlls = F.cross_entropy(logits[..., :-1],
                                           labels[..., 1:],
                                           reduction='none')
                    valid_tokens = attention_mask[..., 1:].bool()
                    self.gen_ppl.update(nlls * valid_tokens, valid_tokens)

    @torch.no_grad()
    def record_tc(self, noise_index, sample):
        self.tc.update(noise_index, sample)

	import numpy as np
	from math import log
	from collections import defaultdict
	import math
	import os
	import typing

	import torch
	import torch.nn.functional as F
	import torchmetrics
	import transformers
	from transformers import logging
	logging.set_verbosity_error()

	# TODO: for tc

	LOG2 = math.log(2)

	# For take the total correlation between x0, x1, and y


	class TC():
	def __init__(self):
	self.x0 = []
	self.x1 = torch.tensor([])

	def update(self, x0, x1):
	# note that x0, x1, y are all torch tensors
	assert len(x0) == len(x1), "All must have same length"

	# it will be a [8, 1024]
	if self.x1.device != x1.device:
	self.x1 = self.x1.to(x1.device)
	self.x0 += x0
	self.x1 = torch.cat((self.x1, x1), dim=0)

	def entropy_from_counts(self, counts, base=2):
	"""
	Compute the empirical entropy H(X) from a dictionary of counts,
	using either base-2 (bits) or base-e (nats).
	"""
	total = sum(counts.values())
	if total == 0:
	return 0.0
	entropy = 0.0
	for c in counts.values():
	if c > 0:
	p = c / total
	# Use math.log(p, base) for bits (base=2) or nats (base=np.e)
	entropy -= p * log(p, base)
	return entropy

	def compute_total_correlation_x1(self, x1_group):
	"""
	Given an array x1_group of shape (N, 16, 16), where each row is
	one sample of x1, estimate T(X1) = sum_j H(X1_j) - H(X1)
	by naive frequency counting.

	We treat each of the 256 pixels in the 16x16 image as a separate
	discrete variable. We compute:
	H(X1) via joint frequencies over all 256 pixels
	H(X1_j) via marginal frequencies of each pixel j
	Then T(X1) = sum_j H(X1_j) - H(X1).
	"""
	N = x1_group.shape[0]
	if N < 2:
	# With <2 samples, you cannot really estimate correlation reliably.
	# We return 0.0 by convention or skip it entirely.
	return 0.0

	# Flatten each image into 256-dimensional vector:
	# shape becomes (N, 1024)
	flattened = x1_group.reshape(N, -1)
	d = flattened.shape[1] # should be 1024

	# ---- Joint distribution (all 256 dims) ----#
	# We'll store each flattened row as a tuple, then count frequencies.
	joint_counts = defaultdict(int)
	for row in flattened:
	key = tuple(row) # row is length-1024
	joint_counts[key] += 1
	H_joint = self.entropy_from_counts(joint_counts, base=2)

	# ---- Marginal distributions (one pixel at a time) ----#
	# We'll compute an entropy for each pixel dimension j.
	marginal_entropies = []
	for j in range(d):
	counts_j = defaultdict(int)
	for row in flattened:
	val_j = row[j].long().item()
	counts_j[val_j] += 1
	H_j = self.entropy_from_counts(counts_j, base=2)
	marginal_entropies.append(H_j)

	sum_marginals = sum(marginal_entropies)

	# ---- Total correlation ----#
	T_val = sum_marginals - H_joint
	return T_val, sum_marginals, H_joint

	def compute_conditional_total_correlation_x1_given_x0y(self, x0, x1):
	"""
	1) Group all samples by unique x0[i] pair.
	- Here x0[i] is a [1024] tensor
	2) Collect the corresponding x1[i] arrays for each group.
	3) For each group, estimate T(X1) by naive frequency counting
	(thus approximating T(X1 \| x0=x0_val, y=y_val)).

	Returns a dict: (x0_bytes) -> estimated total correlation in bits.
	"""

	# Safety checks:
	assert len(x0) == len(x1), "All must have same length"
	N = len(x0)

	# Group x1 by unique x0
	groups = defaultdict(list)
	for i in range(N):
	# We need a hashable key for x0[i], which is shape (1028)
	# Convert to bytes, or a tuple if you prefer
	x0_key = x0[i]
	key = x0_key
	groups[key].append(x1[i])

	# Compute total correlation for each group
	results = {}
	marginals = {}
	joints = {}
	for key, x0_x1_tup in groups.items():
	# x1_list is a tensor, shape each (1024)
	# print(len(x1_list))
	# shape = (num_samples_for_this_group, 1024)
	group = torch.stack(x0_x1_tup, dim=0)
	T_val, sum_marginals, H_joint = self.compute_total_correlation_x1(
	group)
	results[key] = T_val
	marginals[key] = sum_marginals
	joints[key] = H_joint

	return results, joints, marginals

	def compute(self):
	"""
	Compute the total correlation of x1 given x0 and y.
	"""
	tc_results, joints, marginals = self.compute_conditional_total_correlation_x1_given_x0y(
	self.x0, self.x1)

	avg_tc = np.mean(list(tc_results.values()))
	avg_joints = np.mean(list(joints.values()))
	avg_marginals = np.mean(list(marginals.values()))

	return avg_tc, avg_joints, avg_marginals


	class NLL(torchmetrics.aggregation.MeanMetric):
	def update(self,
	value: typing.Union[float, torch.Tensor],
	weight: typing.Union[float, torch.Tensor] = 1.0) -> None:
	"""Update state with data.

	Args:
	value: Either a float or tensor containing data.
	Additional tensor dimensions will be flattened
	weight: Either a float or tensor containing weights
	for calculating the average. Shape of weight should
	be able to broadcast with the shape of `value`.
	Default to `1.0` corresponding to simple harmonic
	average.
	"""
	# broadcast weight to value shape
	if not isinstance(value, torch.Tensor):
	value = torch.as_tensor(value, dtype=self.dtype,
	device=self.device)
	if (weight is not None
	and not isinstance(weight, torch.Tensor)):
	weight = torch.as_tensor(weight,
	dtype=self.dtype,
	device=self.device)
	weight = torch.broadcast_to(weight, value.shape)
	value, weight = self._cast_and_nan_check_input(value,
	weight)

	if value.numel() == 0:
	return
	self.mean_value += value.sum()
	self.weight += weight.sum()


	class BPD(NLL):
	def compute(self) -> torch.Tensor:
	"""Computes the bits per dimension.

	Returns:
	bpd
	"""
	return self.mean_value / self.weight / LOG2


	class Perplexity(NLL):
	def compute(self) -> torch.Tensor:
	"""Computes the Perplexity.

	Returns:
	Perplexity
	"""
	return torch.exp(self.mean_value / self.weight)


	class Metrics:
	def __init__(self, gen_ppl_eval_model_name_or_path=None,
	eval_ppl_batch_size=None) -> None:
	metrics = torchmetrics.MetricCollection({
	'nll': NLL(), 'bpd': BPD(), 'ppl': Perplexity()})
	metrics.set_dtype(torch.float64)
	self.train_nlls = metrics.clone(prefix='train/')
	self.train_aux = BPD()
	self.valid_nlls = metrics.clone(prefix='val/')
	self.valid_aux = BPD()
	self.gen_ppl = Perplexity()
	self.sample_entropy = torchmetrics.aggregation.MeanMetric()
	self.unique_token_count = 0
	self.tc = TC()
	self.eval_ppl_batch_size = eval_ppl_batch_size
	self.gen_ppl_eval_model_name_or_path = gen_ppl_eval_model_name_or_path
	self.tokenizer = transformers.AutoTokenizer.\
	from_pretrained(gen_ppl_eval_model_name_or_path)
	if self.tokenizer.pad_token is None:
	self.tokenizer.pad_token = self.tokenizer.eos_token
	self.tokenizer.pad_token_id = self.tokenizer.eos_token_id

	def to(self, args, *kwargs):
	self.gen_ppl = self.gen_ppl.to(args, *kwargs)
	self.sample_entropy = self.sample_entropy.to(args, *kwargs)
	self.train_nlls = self.train_nlls.to(args, *kwargs)
	self.train_aux = self.train_aux.to(args, *kwargs)
	self.valid_nlls = self.valid_nlls.to(args, *kwargs)
	self.valid_aux = self.valid_aux.to(args, *kwargs)

	def reset(self):
	self.gen_ppl.reset()
	self.sample_entropy.reset()
	self.train_nlls.reset()
	self.train_aux.reset()
	self.valid_nlls.reset()
	self.valid_aux.reset()

	def update_train(self, nll, aux_loss, num_tokens):
	self.train_nlls.update(nll, num_tokens)
	self.train_aux.update(aux_loss, num_tokens)

	def update_valid(self, nll, aux_loss, num_tokens):
	self.valid_nlls.update(nll, num_tokens)
	self.valid_aux.update(aux_loss, num_tokens)

	@torch.no_grad()
	def _eval_retokenize(self, text_samples, max_length,
	device):
	"""Retokenizes samples for the eval model.

	Args:
	text_samples: List of sentences generated by the model.
	Returns:
	samples: Samples re-tokenized for the eval model
	attn_mask: Attention mask for the eval model
	eval_context_size: Size of the context for the eval model
	"""
	if 'llama2' in self.gen_ppl_eval_model_name_or_path.lower():
	tokenizer_kwargs = {
	'text_samples': text_samples,
	'return_tensors': 'pt',
	'return_token_type_ids': False,
	'return_attention_mask': True,
	'truncation': True,
	'padding': True,
	'max_length': max_length,
	}
	eval_context_size = 4096
	elif 'llama3' in self.gen_ppl_eval_model_name_or_path.lower():
	tokenizer_kwargs = {
	'text': text_samples,
	'return_tensors': 'pt',
	'padding': True,
	}
	eval_context_size = 8192
	else:
	tokenizer_kwargs = {
	'return_tensors': 'pt',
	'return_token_type_ids': False,
	'return_attention_mask': True,
	'truncation': True,
	'padding': True,
	'max_length': max_length,
	}
	eval_context_size = 1024
	samples = self.tokenizer(text_samples,
	**tokenizer_kwargs)
	attn_mask = samples['attention_mask']
	samples = samples['input_ids']
	if 'llama' not in self.gen_ppl_eval_model_name_or_path.lower():
	attn_mask = attn_mask.to(device)
	samples = samples.to(device)
	return samples, attn_mask, eval_context_size

	@torch.no_grad()
	def record_entropy(self, tokens):
	for sample in tokens:
	_, counts = torch.unique(
	sample, return_counts=True, sorted=False)
	entropy = torch.special.entr(
	counts.float() / counts.sum()).sum().item()
	self.sample_entropy.update(entropy)

	@torch.no_grad()
	def record_unique_tokens(self, samples):
	"""Record the count of unique tokens across all samples.

	Args:
	samples: torch.Tensor of shape (batch_size, seq_length)

	Returns:
	int: Number of unique tokens in the samples
	"""
	unique_tokens = torch.unique(samples.flatten())
	self.unique_token_count = len(unique_tokens)
	return self.unique_token_count

	def reset_unique_tokens(self):
	"""Reset unique token count."""
	self.unique_token_count = 0

	@torch.no_grad()
	def record_generative_perplexity(
	self,
	text_samples: typing.List[str],
	max_length: int,
	retokenize: bool = True,
	device='cuda') -> None:

	os.environ['TOKENIZERS_PARALLELISM'] = 'false'
	if 'llama' not in self.gen_ppl_eval_model_name_or_path:
	eval_model = transformers.AutoModelForCausalLM.from_pretrained(
	self.gen_ppl_eval_model_name_or_path).eval()
	eval_model = eval_model.to(device)
	# Re-tokenize using eval model's tokenizer
	if retokenize:
	(samples, attn_mask,
	eval_context_size) = self._eval_retokenize(
	text_samples, max_length=max_length, device=device)
	else:
	samples = text_samples
	attn_mask = torch.ones(samples.shape).to(device)
	eval_context_size = samples.shape[-1]
	batch_size = min(self.eval_ppl_batch_size,
	samples.shape[0])
	num_batches = samples.shape[0] // batch_size
	for i in range(num_batches):
	_samples = torch.split(
	samples[i * batch_size: (i + 1) * batch_size],
	eval_context_size,
	dim=-1)
	_attn_mask = torch.split(
	attn_mask[i * batch_size: (i + 1) * batch_size],
	eval_context_size,
	dim=-1)
	for (sample_chunk, attn_mask_chunk) in zip(_samples,
	_attn_mask):
	logits = eval_model(sample_chunk.to(device),
	attention_mask=attn_mask_chunk.to(device))
	logits = logits[0].transpose(-1, -2)
	nlls = F.cross_entropy(logits[..., :-1],
	sample_chunk[..., 1:],
	reduction='none')
	first_eos = (
	sample_chunk
	== self.tokenizer.eos_token_id).cumsum(-1) == 1
	token_mask = sample_chunk != self.tokenizer.eos_token_id
	valid_tokens = first_eos[..., 1:] + token_mask[..., 1:]
	self.gen_ppl.update(nlls * valid_tokens, valid_tokens)
	else:
	eval_model = transformers.AutoModelForCausalLM.from_pretrained(
	self.gen_ppl_eval_model_name_or_path,
	torch_dtype=torch.bfloat16).eval()
	eval_model = eval_model.to(device)
	# Re-tokenize using eval model's tokenizer
	tokenizer_llama = transformers.AutoTokenizer.from_pretrained(
	self.gen_ppl_eval_model_name_or_path)
	if tokenizer_llama.pad_token is None:
	tokenizer_llama.pad_token = tokenizer_llama.eos_token
	tokenizer_llama.pad_token_id = tokenizer_llama.eos_token_id
	tokenizer_gpt = transformers.AutoTokenizer.from_pretrained(
	'gpt2')
	if tokenizer_gpt.pad_token is None:
	tokenizer_gpt.pad_token = tokenizer_gpt.eos_token
	tokenizer_gpt.pad_token_id = tokenizer_gpt.eos_token_id
	num_samples = len(text_samples)
	batch_size = min(16, num_samples)
	num_batches = num_samples // batch_size

	with torch.inference_mode():
	# 1. divide into batches of 16
	# 2. encode each batch
	# 3. eval each batch of 16
	for i in range(num_batches):
	batch_text_samples = text_samples[i *
	batch_size:(i+1)*batch_size]
	encoded_inputs = tokenizer_llama(
	batch_text_samples,
	return_tensors="pt",
	padding=True,
	)

	input_ids = encoded_inputs['input_ids'].to(device)
	attention_mask = encoded_inputs['attention_mask'].to(
	device)

	labels = input_ids.clone()
	labels[labels == tokenizer_llama.pad_token_id] = 50000
	labels = labels.to(device)

	outputs = eval_model(
	input_ids=input_ids, attention_mask=attention_mask, labels=labels)

	llama_logits = outputs.logits

	logits = llama_logits.transpose(-1, -2)
	nlls = F.cross_entropy(logits[..., :-1],
	labels[..., 1:],
	reduction='none')
	valid_tokens = attention_mask[..., 1:].bool()
	self.gen_ppl.update(nlls * valid_tokens, valid_tokens)

	@torch.no_grad()
	def record_tc(self, noise_index, sample):
	self.tc.update(noise_index, sample)