



import argparse
import csv
import json
import math
import time

import numpy as np
import torch
import torch.optim as optim
import torch.utils.data as data
from nltk.tokenize.treebank import TreebankWordDetokenizer
from pplm_classification_head import ClassificationHead
from torch import nn
from torchtext import data as torchtext_data
from torchtext import datasets
from tqdm import tqdm, trange

from transformers import GPTaLMHeadModel, GPTaTokenizer


torch.manual_seed(0)
np.random.seed(0)
lowercase_  =   1e-10
lowercase_  =   """This is incredible! I love it, this is the best chicken I have ever had."""
lowercase_  =   100
class  _snake_case	(    nn.Module):




				def __init__(  self    :  Optional[int],    __lowercase    :  str,    __lowercase    :  str="gpt2-medium",    __lowercase    :  str=False,    __lowercase    :  Tuple="cpu"  ):
							super().__init__()
							lowercase__         = GPTaTokenizer.from_pretrained(__a  )
							lowercase__         = GPTaLMHeadModel.from_pretrained(__a  )
							lowercase__         = self.encoder.transformer.config.hidden_size
							lowercase__         = ClassificationHead(class_size=__a,    embed_size=self.embed_size  )
							lowercase__         = cached_mode
							lowercase__         = device




				def 							A__					(  self    :  Dict  ):
							return self.classifier_head




				def 							A__					(  self    :  int  ):
							for param in self.encoder.parameters():
										lowercase__         = False
							self.classifier_head.train()




				def 							A__					(  self    :  str,    __lowercase    :  Optional[int]  ):
							lowercase__         = x.ne(0  ).unsqueeze(2  ).repeat(1,    1,    self.embed_size  ).float().to(self.device  ).detach()
							lowercase__         = self.encoder.transformer(__a  )["last_hidden_state"]
							lowercase__         = hidden * mask
							lowercase__         = torch.sum(__a,    dim=1  ) / (torch.sum(__a,    dim=1  ).detach() + EPSILON)
							return avg_hidden




				def 							A__					(  self    :  str,    __lowercase    :  Optional[Any]  ):
							if self.cached_mode:
										lowercase__         = x.to(self.device  )
							else:
										lowercase__         = self.avg_representation(x.to(self.device  )  )

							lowercase__         = self.classifier_head(__a  )
							lowercase__         = nn.functional.log_softmax(__a,    dim=-1  )

							return probs
class  _snake_case	(    data.Dataset):




				def __init__(  self    :  Optional[Any],    __lowercase    :  str,    __lowercase    :  Optional[Any]  ):
							lowercase__         = X
							lowercase__         = y




				def __len__(  self    :  Tuple  ):
							return len(self.X  )




				def __getitem__(  self    :  Tuple,    __lowercase    :  List[str]  ):
							lowercase__         = {}
							lowercase__         = self.X[index]
							lowercase__         = self.y[index]
							return data



def __lowerCAmelCase       (   SCREAMING_SNAKE_CASE_     ):

			def pad_sequences(SCREAMING_SNAKE_CASE_     ):
						lowercase__         = [len(lowercase_     ) for seq in sequences]

						lowercase__         = torch.zeros(len(lowercase_     )  ,     max(lowercase_     )     ).long()  # padding value = 0

						for i, seq in enumerate(lowercase_     ):
									lowercase__         = lengths[i]
									lowercase__         = seq[:end]

						return padded_sequences, lengths

			lowercase__         = {}
			for key in data[0].keys():
						lowercase__         = [d[key] for d in data]

			lowercase__         = pad_sequences(item_info["X"]     )
			lowercase__         = torch.tensor(item_info["y"]  ,     dtype=torch.long     )

			return x_batch, y_batch



def __lowerCAmelCase       (   SCREAMING_SNAKE_CASE_     ):
			lowercase__         = {}
			for key in data[0].keys():
						lowercase__         = [d[key] for d in data]

			lowercase__         = torch.cat(item_info["X"]  ,     0     )
			lowercase__         = torch.tensor(item_info["y"]  ,     dtype=torch.long     )

			return x_batch, y_batch



def __lowerCAmelCase       (   SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_=0  ,     SCREAMING_SNAKE_CASE_=10  ,     SCREAMING_SNAKE_CASE_="cpu"     ):
			lowercase__         = 0
			discriminator.train_custom()
			for batch_idx, (input_t, target_t) in enumerate(lowercase_     ):
						lowercase__         = input_t.to(lowercase_     ), target_t.to(lowercase_     )

						optimizer.zero_grad()

						lowercase__         = discriminator(lowercase_     )
						lowercase__         = nn.functional.nll_loss(lowercase_  ,     lowercase_     )
						loss.backward(retain_graph=lowercase_     )
						optimizer.step()

						samples_so_far += len(lowercase_     )

						if batch_idx % log_interval == 0:
									print(
									    "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
									        epoch + 1  ,     lowercase_  ,     len(data_loader.dataset     )  ,     100 * samples_so_far / len(data_loader.dataset     )  ,     loss.item()  ,     )     )



def __lowerCAmelCase       (   SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_="cpu"     ):
			discriminator.eval()
			lowercase__         = 0
			lowercase__         = 0
			with torch.no_grad():
						for input_t, target_t in data_loader:
									lowercase__         = input_t.to(lowercase_     ), target_t.to(lowercase_     )
									lowercase__         = discriminator(lowercase_     )
									# sum up batch loss
									test_loss += nn.functional.nll_loss(lowercase_  ,     lowercase_  ,     reduction="sum"     ).item()
									# get the index of the max log-probability
									lowercase__         = output_t.argmax(dim=1  ,     keepdim=lowercase_     )
									correct += pred_t.eq(target_t.view_as(lowercase_     )     ).sum().item()

			test_loss /= len(data_loader.dataset     )

			print(
			    "Performance on test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)".format(
			        lowercase_  ,     lowercase_  ,     len(data_loader.dataset     )  ,     100.0 * correct / len(data_loader.dataset     )     )     )



def __lowerCAmelCase       (   SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_=False  ,     SCREAMING_SNAKE_CASE_="cpu"     ):
			lowercase__         = model.tokenizer.encode(lowercase_     )
			lowercase__         = torch.tensor([input_t]  ,     dtype=torch.long  ,     device=lowercase_     )
			if cached:
						lowercase__         = model.avg_representation(lowercase_     )

			lowercase__         = model(lowercase_     ).data.cpu().numpy().flatten().tolist()
			print("Input sentence:"  ,     lowercase_     )
			print(
			    "Predictions:"  ,     ", ".join("{}: {:.4f}".format(lowercase_  ,     math.exp(lowercase_     )     ) for c, log_prob in zip(lowercase_  ,     lowercase_     )     )  ,     )



def __lowerCAmelCase       (   SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_=False  ,     SCREAMING_SNAKE_CASE_="cpu"     ):
			lowercase__         = torch.utils.data.DataLoader(dataset=lowercase_  ,     batch_size=lowercase_  ,     collate_fn=lowercase_     )

			lowercase__         = []
			lowercase__         = []
			for batch_idx, (x, y) in enumerate(tqdm(lowercase_  ,     ascii=lowercase_     )     ):
						with torch.no_grad():
									lowercase__         = x.to(lowercase_     )
									lowercase__         = discriminator.avg_representation(lowercase_     ).cpu().detach()
									lowercase__         = torch.unbind(avg_rep.unsqueeze(1     )     )
									xs += avg_rep_list
									ys += y.cpu().numpy().tolist()

			lowercase__         = torch.utils.data.DataLoader(
			    dataset=Dataset(lowercase_  ,     lowercase_     )  ,     batch_size=lowercase_  ,     shuffle=lowercase_  ,     collate_fn=lowercase_     )

			return data_loader



def __lowerCAmelCase       (   SCREAMING_SNAKE_CASE_  ,     SCREAMING_SNAKE_CASE_=None  ,     SCREAMING_SNAKE_CASE_="gpt2-medium"  ,     SCREAMING_SNAKE_CASE_=10  ,     SCREAMING_SNAKE_CASE_=64  ,     SCREAMING_SNAKE_CASE_=10  ,     SCREAMING_SNAKE_CASE_=False  ,     SCREAMING_SNAKE_CASE_=False  ,     SCREAMING_SNAKE_CASE_=False  ,     ):
			lowercase__         = "cuda" if torch.cuda.is_available() and not no_cuda else "cpu"

			print("Preprocessing {} dataset...".format(lowercase_     )     )
			lowercase__         = time.time()

			if dataset == "SST":
						lowercase__         = ["positive", "negative", "very positive", "very negative", "neutral"]
						lowercase__         = {c: i for i, c in enumerate(lowercase_     )}

						lowercase__         = Discriminator(
						    class_size=len(lowercase_     )  ,     pretrained_model=lowercase_  ,     cached_mode=lowercase_  ,     device=lowercase_     ).to(lowercase_     )

						lowercase__         = torchtext_data.Field()
						lowercase__         = torchtext_data.Field(sequential=lowercase_     )
						lowercase__         = datasets.SST.splits(
						    lowercase_  ,     lowercase_  ,     fine_grained=lowercase_  ,     train_subtrees=lowercase_  ,     )

						lowercase__         = []
						lowercase__         = []
						for i in trange(len(lowercase_     )  ,     ascii=lowercase_     ):
									lowercase__         = TreebankWordDetokenizer().detokenize(vars(train_data[i]     )["text"]     )
									lowercase__         = discriminator.tokenizer.encode(lowercase_     )
									lowercase__         = torch.tensor([5_0256] + seq  ,     device=lowercase_  ,     dtype=torch.long     )
									x.append(lowercase_     )
									y.append(classaidx[vars(train_data[i]     )["label"]]     )
						lowercase__         = Dataset(lowercase_  ,     lowercase_     )

						lowercase__         = []
						lowercase__         = []
						for i in trange(len(lowercase_     )  ,     ascii=lowercase_     ):
									lowercase__         = TreebankWordDetokenizer().detokenize(vars(test_data[i]     )["text"]     )
									lowercase__         = discriminator.tokenizer.encode(lowercase_     )
									lowercase__         = torch.tensor([5_0256] + seq  ,     device=lowercase_  ,     dtype=torch.long     )
									test_x.append(lowercase_     )
									test_y.append(classaidx[vars(test_data[i]     )["label"]]     )
						lowercase__         = Dataset(lowercase_  ,     lowercase_     )

						lowercase__         = {
						    "class_size": len(lowercase_     ),
						    "embed_size": discriminator.embed_size,
						    "pretrained_model": pretrained_model,
						    "class_vocab": classaidx,
						    "default_class": 2,
						}

			elif dataset == "clickbait":
						lowercase__         = ["non_clickbait", "clickbait"]
						lowercase__         = {c: i for i, c in enumerate(lowercase_     )}

						lowercase__         = Discriminator(
						    class_size=len(lowercase_     )  ,     pretrained_model=lowercase_  ,     cached_mode=lowercase_  ,     device=lowercase_     ).to(lowercase_     )

						with open("datasets/clickbait/clickbait_train_prefix.txt"     ) as f:
									lowercase__         = []
									for i, line in enumerate(lowercase_     ):
												try:
															data.append(eval(lowercase_     )     )
												except Exception:
															print("Error evaluating line {}: {}".format(lowercase_  ,     lowercase_     )     )
															continue
						lowercase__         = []
						lowercase__         = []
						with open("datasets/clickbait/clickbait_train_prefix.txt"     ) as f:
									for i, line in enumerate(tqdm(lowercase_  ,     ascii=lowercase_     )     ):
												try:
															lowercase__         = eval(lowercase_     )
															lowercase__         = discriminator.tokenizer.encode(d["text"]     )

															if len(lowercase_     ) < max_length_seq:
																		lowercase__         = torch.tensor([5_0256] + seq  ,     device=lowercase_  ,     dtype=torch.long     )
															else:
																		print("Line {} is longer than maximum length {}".format(lowercase_  ,     lowercase_     )     )
																		continue
															x.append(lowercase_     )
															y.append(d["label"]     )
												except Exception:
															print("Error evaluating / tokenizing line {}, skipping it".format(lowercase_     )     )
															pass

						lowercase__         = Dataset(lowercase_  ,     lowercase_     )
						lowercase__         = int(0.9 * len(lowercase_     )     )
						lowercase__         = len(lowercase_     ) - train_size
						lowercase__         = torch.utils.data.random_split(lowercase_  ,     [train_size, test_size]     )

						lowercase__         = {
						    "class_size": len(lowercase_     ),
						    "embed_size": discriminator.embed_size,
						    "pretrained_model": pretrained_model,
						    "class_vocab": classaidx,
						    "default_class": 1,
						}

			elif dataset == "toxic":
						lowercase__         = ["non_toxic", "toxic"]
						lowercase__         = {c: i for i, c in enumerate(lowercase_     )}

						lowercase__         = Discriminator(
						    class_size=len(lowercase_     )  ,     pretrained_model=lowercase_  ,     cached_mode=lowercase_  ,     device=lowercase_     ).to(lowercase_     )

						lowercase__         = []
						lowercase__         = []
						with open("datasets/toxic/toxic_train.txt"     ) as f:
									for i, line in enumerate(tqdm(lowercase_  ,     ascii=lowercase_     )     ):
												try:
															lowercase__         = eval(lowercase_     )
															lowercase__         = discriminator.tokenizer.encode(d["text"]     )

															if len(lowercase_     ) < max_length_seq:
																		lowercase__         = torch.tensor([5_0256] + seq  ,     device=lowercase_  ,     dtype=torch.long     )
															else:
																		print("Line {} is longer than maximum length {}".format(lowercase_  ,     lowercase_     )     )
																		continue
															x.append(lowercase_     )
															y.append(int(np.sum(d["label"]     ) > 0     )     )
												except Exception:
															print("Error evaluating / tokenizing line {}, skipping it".format(lowercase_     )     )
															pass

						lowercase__         = Dataset(lowercase_  ,     lowercase_     )
						lowercase__         = int(0.9 * len(lowercase_     )     )
						lowercase__         = len(lowercase_     ) - train_size
						lowercase__         = torch.utils.data.random_split(lowercase_  ,     [train_size, test_size]     )

						lowercase__         = {
						    "class_size": len(lowercase_     ),
						    "embed_size": discriminator.embed_size,
						    "pretrained_model": pretrained_model,
						    "class_vocab": classaidx,
						    "default_class": 0,
						}

			else:  # if dataset == "generic":
						# This assumes the input dataset is a TSV with the following structure:
						# class \t text

						if dataset_fp is None:
									raise ValueError("When generic dataset is selected, dataset_fp needs to be specified aswell."     )

						lowercase__         = set()
						with open(lowercase_     ) as f:
									lowercase__         = csv.reader(lowercase_  ,     delimiter="\t"     )
									for row in tqdm(lowercase_  ,     ascii=lowercase_     ):
												if row:
															classes.add(row[0]     )

						lowercase__         = sorted(lowercase_     )
						lowercase__         = {c: i for i, c in enumerate(lowercase_     )}

						lowercase__         = Discriminator(
						    class_size=len(lowercase_     )  ,     pretrained_model=lowercase_  ,     cached_mode=lowercase_  ,     device=lowercase_     ).to(lowercase_     )

						lowercase__         = []
						lowercase__         = []
						with open(lowercase_     ) as f:
									lowercase__         = csv.reader(lowercase_  ,     delimiter="\t"     )
									for i, row in enumerate(tqdm(lowercase_  ,     ascii=lowercase_     )     ):
												if row:
															lowercase__         = row[0]
															lowercase__         = row[1]

															try:
																		lowercase__         = discriminator.tokenizer.encode(lowercase_     )
																		if len(lowercase_     ) < max_length_seq:
																					lowercase__         = torch.tensor([5_0256] + seq  ,     device=lowercase_  ,     dtype=torch.long     )

																		else:
																					print("Line {} is longer than maximum length {}".format(lowercase_  ,     lowercase_     )     )
																					continue

																		x.append(lowercase_     )
																		y.append(classaidx[label]     )

															except Exception:
																		print("Error tokenizing line {}, skipping it".format(lowercase_     )     )
																		pass

						lowercase__         = Dataset(lowercase_  ,     lowercase_     )
						lowercase__         = int(0.9 * len(lowercase_     )     )
						lowercase__         = len(lowercase_     ) - train_size
						lowercase__         = torch.utils.data.random_split(lowercase_  ,     [train_size, test_size]     )

						lowercase__         = {
						    "class_size": len(lowercase_     ),
						    "embed_size": discriminator.embed_size,
						    "pretrained_model": pretrained_model,
						    "class_vocab": classaidx,
						    "default_class": 0,
						}

			lowercase__         = time.time()
			print("Preprocessed {} data points".format(len(lowercase_     ) + len(lowercase_     )     )     )
			print("Data preprocessing took: {:.3f}s".format(end - start     )     )

			if cached:
						print("Building representation cache..."     )

						lowercase__         = time.time()

						lowercase__         = get_cached_data_loader(lowercase_  ,     lowercase_  ,     lowercase_  ,     shuffle=lowercase_  ,     device=lowercase_     )

						lowercase__         = get_cached_data_loader(lowercase_  ,     lowercase_  ,     lowercase_  ,     device=lowercase_     )

						lowercase__         = time.time()
						print("Building representation cache took: {:.3f}s".format(end - start     )     )

			else:
						lowercase__         = torch.utils.data.DataLoader(
						    dataset=lowercase_  ,     batch_size=lowercase_  ,     shuffle=lowercase_  ,     collate_fn=lowercase_     )
						lowercase__         = torch.utils.data.DataLoader(dataset=lowercase_  ,     batch_size=lowercase_  ,     collate_fn=lowercase_     )

			if save_model:
						with open("{}_classifier_head_meta.json".format(lowercase_     )  ,     "w"     ) as meta_file:
									json.dump(lowercase_  ,     lowercase_     )

			lowercase__         = optim.Adam(discriminator.parameters()  ,     lr=0.0001     )

			for epoch in range(lowercase_     ):
						lowercase__         = time.time()
						print("\nEpoch"  ,     epoch + 1     )

						train_epoch(
						    discriminator=lowercase_  ,     data_loader=lowercase_  ,     optimizer=lowercase_  ,     epoch=lowercase_  ,     log_interval=lowercase_  ,     device=lowercase_  ,     )
						evaluate_performance(data_loader=lowercase_  ,     discriminator=lowercase_  ,     device=lowercase_     )

						lowercase__         = time.time()
						print("Epoch took: {:.3f}s".format(end - start     )     )

						print("\nExample prediction"     )
						predict(lowercase_  ,     lowercase_  ,     lowercase_  ,     cached=lowercase_  ,     device=lowercase_     )

						if save_model:
									# torch.save(discriminator.state_dict(),
									#           "{}_discriminator_{}.pt".format(
									#               args.dataset, epoch + 1
									#               ))
									torch.save(
									    discriminator.get_classifier().state_dict()  ,     "{}_classifier_head_epoch_{}.pt".format(lowercase_  ,     epoch + 1     )  ,     )


if __name__ == "__main__":
							lowercase_  =   argparse.ArgumentParser(description="""Train a discriminator on top of GPT-2 representations""")
							parser.add_argument(
							    """--dataset""",
							    type=str,
							    default="""SST""",
							    choices=("""SST""", """clickbait""", """toxic""", """generic"""),
							    help=(
							        """dataset to train the discriminator on."""
							        """In case of generic, the dataset is expected"""
							        """to be a TSBV file with structure: class \\t text"""
							    ),
							)
							parser.add_argument(
							    """--dataset_fp""",
							    type=str,
							    default="""""",
							    help="""File path of the dataset to use. Needed only in case of generic datadset""",
							)
							parser.add_argument(
							    """--pretrained_model""", type=str, default="""gpt2-medium""", help="""Pretrained model to use as encoder"""
							)
							parser.add_argument("""--epochs""", type=int, default=10, metavar="""N""", help="""Number of training epochs""")
							parser.add_argument(
							    """--batch_size""", type=int, default=64, metavar="""N""", help="""input batch size for training (default: 64)"""
							)
							parser.add_argument(
							    """--log_interval""",
							    type=int,
							    default=10,
							    metavar="""N""",
							    help="""how many batches to wait before logging training status""",
							)
							parser.add_argument("""--save_model""", action="""store_true""", help="""whether to save the model""")
							parser.add_argument("""--cached""", action="""store_true""", help="""whether to cache the input representations""")
							parser.add_argument("""--no_cuda""", action="""store_true""", help="""use to turn off cuda""")
							lowercase_  =   parser.parse_args()

							train_discriminator(**(vars(args)))











"""simple docstring"""


from typing import TYPE_CHECKING

from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available


lowerCamelCase__      =      {
    "configuration_utils": ["GenerationConfig"],
    "streamers": ["TextIteratorStreamer", "TextStreamer"],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase__      =      [
      "Constraint",
      "ConstraintListState",
      "DisjunctiveConstraint",
      "PhrasalConstraint",
  ]
  lowerCamelCase__      =      [
      "BeamHypotheses",
      "BeamScorer",
      "BeamSearchScorer",
      "ConstrainedBeamSearchScorer",
  ]
  lowerCamelCase__      =      [
      "EpsilonLogitsWarper",
      "EtaLogitsWarper",
      "ForcedBOSTokenLogitsProcessor",
      "ForcedEOSTokenLogitsProcessor",
      "HammingDiversityLogitsProcessor",
      "InfNanRemoveLogitsProcessor",
      "LogitsProcessor",
      "LogitsProcessorList",
      "LogitsWarper",
      "MinLengthLogitsProcessor",
      "MinNewTokensLengthLogitsProcessor",
      "NoBadWordsLogitsProcessor",
      "NoRepeatNGramLogitsProcessor",
      "PrefixConstrainedLogitsProcessor",
      "RepetitionPenaltyLogitsProcessor",
      "SequenceBiasLogitsProcessor",
      "EncoderRepetitionPenaltyLogitsProcessor",
      "TemperatureLogitsWarper",
      "TopKLogitsWarper",
      "TopPLogitsWarper",
      "TypicalLogitsWarper",
      "EncoderNoRepeatNGramLogitsProcessor",
      "ExponentialDecayLengthPenalty",
      "LogitNormalization",
  ]
  lowerCamelCase__      =      [
      "MaxNewTokensCriteria",
      "MaxLengthCriteria",
      "MaxTimeCriteria",
      "StoppingCriteria",
      "StoppingCriteriaList",
      "validate_stopping_criteria",
  ]
  lowerCamelCase__      =      [
      "GenerationMixin",
      "top_k_top_p_filtering",
      "GreedySearchEncoderDecoderOutput",
      "GreedySearchDecoderOnlyOutput",
      "SampleEncoderDecoderOutput",
      "SampleDecoderOnlyOutput",
      "BeamSearchEncoderDecoderOutput",
      "BeamSearchDecoderOnlyOutput",
      "BeamSampleEncoderDecoderOutput",
      "BeamSampleDecoderOnlyOutput",
      "ContrastiveSearchEncoderDecoderOutput",
      "ContrastiveSearchDecoderOnlyOutput",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase__      =      [
      "TFForcedBOSTokenLogitsProcessor",
      "TFForcedEOSTokenLogitsProcessor",
      "TFLogitsProcessor",
      "TFLogitsProcessorList",
      "TFLogitsWarper",
      "TFMinLengthLogitsProcessor",
      "TFNoBadWordsLogitsProcessor",
      "TFNoRepeatNGramLogitsProcessor",
      "TFRepetitionPenaltyLogitsProcessor",
      "TFTemperatureLogitsWarper",
      "TFTopKLogitsWarper",
      "TFTopPLogitsWarper",
      "TFForceTokensLogitsProcessor",
      "TFSuppressTokensAtBeginLogitsProcessor",
      "TFSuppressTokensLogitsProcessor",
  ]
  lowerCamelCase__      =      [
      "TFGenerationMixin",
      "tf_top_k_top_p_filtering",
      "TFGreedySearchDecoderOnlyOutput",
      "TFGreedySearchEncoderDecoderOutput",
      "TFSampleEncoderDecoderOutput",
      "TFSampleDecoderOnlyOutput",
      "TFBeamSearchEncoderDecoderOutput",
      "TFBeamSearchDecoderOnlyOutput",
      "TFBeamSampleEncoderDecoderOutput",
      "TFBeamSampleDecoderOnlyOutput",
      "TFContrastiveSearchEncoderDecoderOutput",
      "TFContrastiveSearchDecoderOnlyOutput",
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase__      =      [
      "FlaxForcedBOSTokenLogitsProcessor",
      "FlaxForcedEOSTokenLogitsProcessor",
      "FlaxLogitsProcessor",
      "FlaxLogitsProcessorList",
      "FlaxLogitsWarper",
      "FlaxMinLengthLogitsProcessor",
      "FlaxTemperatureLogitsWarper",
      "FlaxTopKLogitsWarper",
      "FlaxTopPLogitsWarper",
  ]
  lowerCamelCase__      =      [
      "FlaxGenerationMixin",
      "FlaxGreedySearchOutput",
      "FlaxSampleOutput",
      "FlaxBeamSearchOutput",
  ]

if TYPE_CHECKING:
  from .configuration_utils import GenerationConfig
  from .streamers import TextIteratorStreamer, TextStreamer

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
    from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
    from .logits_process import (
        EncoderNoRepeatNGramLogitsProcessor,
        EncoderRepetitionPenaltyLogitsProcessor,
        EpsilonLogitsWarper,
        EtaLogitsWarper,
        ExponentialDecayLengthPenalty,
        ForcedBOSTokenLogitsProcessor,
        ForcedEOSTokenLogitsProcessor,
        HammingDiversityLogitsProcessor,
        InfNanRemoveLogitsProcessor,
        LogitNormalization,
        LogitsProcessor,
        LogitsProcessorList,
        LogitsWarper,
        MinLengthLogitsProcessor,
        MinNewTokensLengthLogitsProcessor,
        NoBadWordsLogitsProcessor,
        NoRepeatNGramLogitsProcessor,
        PrefixConstrainedLogitsProcessor,
        RepetitionPenaltyLogitsProcessor,
        SequenceBiasLogitsProcessor,
        TemperatureLogitsWarper,
        TopKLogitsWarper,
        TopPLogitsWarper,
        TypicalLogitsWarper,
    )
    from .stopping_criteria import (
        MaxLengthCriteria,
        MaxNewTokensCriteria,
        MaxTimeCriteria,
        StoppingCriteria,
        StoppingCriteriaList,
        validate_stopping_criteria,
    )
    from .utils import (
        BeamSampleDecoderOnlyOutput,
        BeamSampleEncoderDecoderOutput,
        BeamSearchDecoderOnlyOutput,
        BeamSearchEncoderDecoderOutput,
        ContrastiveSearchDecoderOnlyOutput,
        ContrastiveSearchEncoderDecoderOutput,
        GenerationMixin,
        GreedySearchDecoderOnlyOutput,
        GreedySearchEncoderDecoderOutput,
        SampleDecoderOnlyOutput,
        SampleEncoderDecoderOutput,
        top_k_top_p_filtering,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tf_logits_process import (
        TFForcedBOSTokenLogitsProcessor,
        TFForcedEOSTokenLogitsProcessor,
        TFForceTokensLogitsProcessor,
        TFLogitsProcessor,
        TFLogitsProcessorList,
        TFLogitsWarper,
        TFMinLengthLogitsProcessor,
        TFNoBadWordsLogitsProcessor,
        TFNoRepeatNGramLogitsProcessor,
        TFRepetitionPenaltyLogitsProcessor,
        TFSuppressTokensAtBeginLogitsProcessor,
        TFSuppressTokensLogitsProcessor,
        TFTemperatureLogitsWarper,
        TFTopKLogitsWarper,
        TFTopPLogitsWarper,
    )
    from .tf_utils import (
        TFBeamSampleDecoderOnlyOutput,
        TFBeamSampleEncoderDecoderOutput,
        TFBeamSearchDecoderOnlyOutput,
        TFBeamSearchEncoderDecoderOutput,
        TFContrastiveSearchDecoderOnlyOutput,
        TFContrastiveSearchEncoderDecoderOutput,
        TFGenerationMixin,
        TFGreedySearchDecoderOnlyOutput,
        TFGreedySearchEncoderDecoderOutput,
        TFSampleDecoderOnlyOutput,
        TFSampleEncoderDecoderOutput,
        tf_top_k_top_p_filtering,
    )

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .flax_logits_process import (
        FlaxForcedBOSTokenLogitsProcessor,
        FlaxForcedEOSTokenLogitsProcessor,
        FlaxLogitsProcessor,
        FlaxLogitsProcessorList,
        FlaxLogitsWarper,
        FlaxMinLengthLogitsProcessor,
        FlaxTemperatureLogitsWarper,
        FlaxTopKLogitsWarper,
        FlaxTopPLogitsWarper,
    )
    from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
  import sys

  lowerCamelCase__      =      _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
