



'''simple docstring'''







import argparse
import logging

import numpy as np
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
    datefmt='%m/%d/%Y %H:%M:%S',
    level=logging.INFO,
)
lowercase__  :	List[Any]            =      logging.getLogger(__name__)





def     snake_case__			(lowerCAmelCase_						):
    np.random.seed(args.seed						)
    torch.manual_seed(args.seed						)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed						)





def     snake_case__			():
    UpperCamelCase__	     =      argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path"       ,	default=lowerCAmelCase_       ,	type=lowerCAmelCase_       ,	required=lowerCAmelCase_       ,	)
    parser.add_argument("--prompt"       ,	type=lowerCAmelCase_       ,	default=""						)
    parser.add_argument("--length"       ,	type=lowerCAmelCase_       ,	default=20						)
    parser.add_argument("--stop_token"       ,	type=lowerCAmelCase_       ,	default=lowerCAmelCase_       ,	help="Token at which text generation is stopped"						)
    parser.add_argument(
        "--temperature"       ,	type=lowerCAmelCase_       ,	default=1.0       ,	help="temperature of 1.0 has no effect, lower tend toward greedy sampling"       ,	)
    parser.add_argument(
        "--repetition_penalty"       ,	type=lowerCAmelCase_       ,	default=1.0       ,	help="primarily useful for CTRL model; in that case, use 1.2"						)
    parser.add_argument("--k"       ,	type=lowerCAmelCase_       ,	default=0						)
    parser.add_argument("--penalty_alpha"       ,	type=lowerCAmelCase_       ,	default=0.0						)
    parser.add_argument("--p"       ,	type=lowerCAmelCase_       ,	default=0.9						)

    parser.add_argument("--prefix"       ,	type=lowerCAmelCase_       ,	default=""       ,	help="Text added prior to input."						)
    parser.add_argument("--padding_text"       ,	type=lowerCAmelCase_       ,	default=""       ,	help="Deprecated, the use of `--prefix` is preferred."						)
    parser.add_argument("--xlm_language"       ,	type=lowerCAmelCase_       ,	default=""       ,	help="Optional language when used with the XLM model."						)

    parser.add_argument("--seed"       ,	type=lowerCAmelCase_       ,	default=42       ,	help="random seed for initialization"						)
    parser.add_argument("--no_cuda"       ,	action="store_true"       ,	help="Avoid using CUDA when available"						)
    parser.add_argument(
        "--fp16"       ,	action="store_true"       ,	help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"       ,	)
    UpperCamelCase__	     =      parser.parse_args()

    UpperCamelCase__	     =      torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu"						)
    UpperCamelCase__	     =      0 if args.no_cuda else torch.cuda.device_count()

    logger.warning(F'device: {args.device}, n_gpu: {args.n_gpu}, 16-bits training: {args.fpaa}'						)

    set_seed(lowerCAmelCase_						)

    # Initialize the model and tokenizer
    UpperCamelCase__	     =      AutoTokenizer.from_pretrained(args.model_name_or_path						)
    UpperCamelCase__	     =      AutoModelForCausalLM.from_pretrained(args.model_name_or_path						)

    # tokenizer = GPT2Tokenizer.from_pretrained(args.model_name_or_path)
    # model = OPTForCausalLM.from_pretrained(args.model_name_or_path)
    model.to(args.device						)

    if args.fpaa:
        model.half()

    logger.info(lowerCAmelCase_						)
    UpperCamelCase__	     =      args.prompt if args.prompt else input("Model prompt >>> "						)

    UpperCamelCase__	     =      tokenizer(lowerCAmelCase_       ,	return_tensors="pt"       ,	add_special_tokens=lowerCAmelCase_						)
    UpperCamelCase__	     =      {key: value.to(args.device						) for key, value in inputs.items()}

    UpperCamelCase__	     =      model.generate(
        **lowerCAmelCase_       ,	max_length=args.length + len(inputs["input_ids"][0]						)       ,	penalty_alpha=args.penalty_alpha       ,	top_k=args.k       ,	)

    UpperCamelCase__	     =      []
    for generated_sequence_idx, generated_sequence in enumerate(lowerCAmelCase_						):
        print(F'=== GENERATED SEQUENCE {generated_sequence_idx + 1} ==='						)
        UpperCamelCase__	     =      generated_sequence.tolist()

        # Decode text
        UpperCamelCase__	     =      tokenizer.decode(lowerCAmelCase_       ,	clean_up_tokenization_spaces=lowerCAmelCase_       ,	add_special_tokens=lowerCAmelCase_						)

        # Remove all text after the stop token
        UpperCamelCase__	     =      text[: text.find(args.stop_token						) if args.stop_token else None]

        # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
        UpperCamelCase__	     =      (
            prompt_text + text[len(tokenizer.decode(inputs["input_ids"][0]       ,	clean_up_tokenization_spaces=lowerCAmelCase_						)						) :]
        )

        generated_sequences.append(lowerCAmelCase_						)
        print(lowerCAmelCase_						)

    return generated_sequences


if __name__ == "__main__":
      main()









'''simple docstring'''







import bisect
import copy
from collections import defaultdict

import numpy as np
from torch.utils.data import BatchSampler, Sampler

from utils import logger





def     snake_case__			(lowerCAmelCase_       ,	lowerCAmelCase_						):
    UpperCamelCase__	     =      copy.deepcopy(lowerCAmelCase_						)
    UpperCamelCase__	     =      sorted(lowerCAmelCase_						)
    UpperCamelCase__	     =      [bisect.bisect_right(lowerCAmelCase_       ,	lowerCAmelCase_						) for y in x]
    return quantized





def     snake_case__			(lowerCAmelCase_       ,	lowerCAmelCase_=0						):
    UpperCamelCase__	     =      np.arange(start=3       ,	stop=lowerCAmelCase_       ,	step=4						).tolist() if k > 0 else [10]
    UpperCamelCase__	     =      _quantize(lowerCAmelCase_       ,	lowerCAmelCase_						)
    # count number of elements per group
    UpperCamelCase__	     =      np.unique(lowerCAmelCase_       ,	return_counts=lowerCAmelCase_						)[1]
    UpperCamelCase__	     =      [0] + bins + [np.inf]
    logger.info("Using {} as bins for aspect lengths quantization".format(lowerCAmelCase_						)						)
    logger.info("Count of instances per bin: {}".format(lowerCAmelCase_						)						)
    return groups





class    lowerCamelCase_				(						lowerCAmelCase_	):






   '''simple docstring'''

   def __init__(	self  ,    __UpperCamelCase  ,    __UpperCamelCase  ,    __UpperCamelCase			):
       if not isinstance(__UpperCamelCase  ,    __UpperCamelCase			):
           raise ValueError(
               "sampler should be an instance of torch.utils.data.Sampler, but got sampler={}".format(__UpperCamelCase			)			)
       UpperCamelCase__	     =      sampler
       UpperCamelCase__	     =      group_ids
       UpperCamelCase__	     =      batch_size


   def __iter__(	self			):
       UpperCamelCase__	     =      defaultdict(__UpperCamelCase			)
       UpperCamelCase__	     =      defaultdict(__UpperCamelCase			)

       UpperCamelCase__	     =      0
       for idx in self.sampler:
           UpperCamelCase__	     =      self.group_ids[idx]
           buffer_per_group[group_id].append(__UpperCamelCase			)
           samples_per_group[group_id].append(__UpperCamelCase			)
           if len(buffer_per_group[group_id]			) == self.batch_size:
               yield buffer_per_group[group_id]  # TODO
               num_batches += 1
               del buffer_per_group[group_id]
           assert len(buffer_per_group[group_id]			) < self.batch_size

       # now we have run out of elements that satisfy
       # the group criteria, let's return the remaining
       # elements so that the size of the sampler is
       # deterministic
       UpperCamelCase__	     =      len(self			)
       UpperCamelCase__	     =      expected_num_batches - num_batches
       if num_remaining > 0:
           # for the remaining batches, group the batches by similar lengths
           UpperCamelCase__	     =      []
           for group_id, idxs in sorted(buffer_per_group.items()  ,    key=lambda __UpperCamelCase			: x[0]			):
               batch_idx.extend(__UpperCamelCase			)
               if len(__UpperCamelCase			) >= self.batch_size:
                   yield batch_idx[: self.batch_size]
                   UpperCamelCase__	     =      batch_idx[self.batch_size :]
                   num_remaining -= 1
           if len(__UpperCamelCase			) > 0:
               yield batch_idx
               num_remaining -= 1
       assert num_remaining == 0







   def __len__(	self			):
       return (len(self.sampler			) + self.batch_size - 1) // self.batch_size




