text: import numpy as np
  import torch
  
  
  from modeling_gpt2 import GPT2LMHeadModel
  
  from transformers import (
      GPT2Config,
      GPT2Tokenizer
  )
  
  from nltk import sent_tokenize
  import nltk
  nltk.download('punkt')
  
  code_desired = "true"
  code_undesired = "false"
  model_type = 'gpt2'
  gen_type = "gedi"
  gen_model_name_or_path = "gpt2-xl"
  
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  
  MODEL_CLASSES = {"gpt2": (GPT2Config, GPT2LMHeadModel, GPT2Tokenizer),}
  config_class, model_class, tokenizer_class = MODEL_CLASSES["gpt2"]
  tokenizer = tokenizer_class.from_pretrained(gen_model_name_or_path, do_lower_case=False)
  
  model = model_class.from_pretrained(gen_model_name_or_path, load_in_half_prec=True)
  model = model.to(device)
  model = model.float()
  
  gedi_model_name_or_path = 'sst_5_delib_100_score_model'
  gedi_model = model_class.from_pretrained(gedi_model_name_or_path)
  gedi_model.to(device)
  
  #setting arguments for generation
  #max generation length
  gen_length = 50
  #omega from paper, higher disc_weight means more aggressive topic steering
  disc_weight = 30
  #1 - rho from paper, should be between 0 and 1 higher filter_p means more aggressive topic steering
  filter_p = 0.8
  #tau from paper, preserves tokens that are classified as correct topic
  target_p = 0.8
  #hyperparameter that determines class prior, set to uniform by default
  class_bias = 0
  
  if gen_length>1024:
    length = 1024
  else:
    length = gen_length
  
  def cut_into_sentences(text, do_cleanup=True):
      """
      Cut text into sentences. \n are also regarded as a sentence.
      :param do_cleanup: if True, do cleanups.
      :param text: input text.
      :return: sentences.
      """
      all_sentences = []
      # sentences_raw = text.split("\n")
      sentences_raw = [text.replace("\n", " ")]
      result = []
  
      for item in sentences_raw:
          sentence_in_item = sent_tokenize(item)
          for item2 in sentence_in_item:
              all_sentences.append(item2)
  
      if do_cleanup:
          for item in all_sentences:
              item = item.replace('<|endoftext|>', '')
              if len(item) > 2:
                  result.append(item)
      else:
          result = all_sentences
      return result
  
  
  
  
  def generate_one_sentence(sentence, control, length=50, disc_weight=30, temperature=0.8, gpt3_id=None):
      """
      Generate one sentence based on input data.
      :param sentence: (string) context (prompt) used.
      :param topic: (dict) {topic: weight, topic:weight,...} topic that the sentence need to steer towards.
      :param extra_args: (dict) a dictionary that certain key will trigger additional functionality.
          disc_weight: Set this value to use a different control strength than default.
          get_gen_token_count: Return only how many tokens the generator has generated (for debug only).
      :return: sentence generated, or others if extra_args are specified.
      """
      secondary_code = control
  
      # disc_weight = self.disc_weight
      # if type(extra_args) is dict and 'disc_weight' in extra_args:
      #     disc_weight = extra_args['disc_weight']
  
      if sentence == "":
          print("Prompt is empty! Using a dummy sentence.")
          sentence = "."
  
      # Specify prompt below
      prompt = sentence
  
      # Calculate oroginal input length.
      length_of_prompt = len(sentence)
  
      start_len = 0
      text_ids = tokenizer.encode(prompt)
      length_of_prompt_in_tokens = len(text_ids)
      encoded_prompts = torch.LongTensor(text_ids).unsqueeze(0).to(device)
  
      if type(control) is str:
          multi_code = tokenizer.encode(secondary_code)
      elif type(control) is dict:
          multi_code = {}
          for item in secondary_code:
              encoded = tokenizer.encode(item)[0]  # only take the first one
              multi_code[encoded] = secondary_code[item]
      else:
          raise NotImplementedError("topic data type of %s not supported... Supported: (str,dict)" % type(control))
  
      # If 1, generate sentences towards a specific topic.
      attr_class = 1
  
      if int(control)!=-1:
        if gpt3_id is None:
          generated_sequence = model.generate(input_ids=encoded_prompts,
                                                    pad_lens=None,
                                                    max_length=length + length_of_prompt_in_tokens,
                                                    top_k=None,
                                                    top_p=None,
                                                    repetition_penalty=1.2,
                                                    rep_penalty_scale=10,
                                                    eos_token_ids=tokenizer.eos_token_id,
                                                    pad_token_id=0,  # self.tokenizer.eos_token_id,
                                                    do_sample=True,
                                                    temperature = temperature,
                                                    penalize_cond=True,
                                                    gedi_model=gedi_model,
                                                    tokenizer=tokenizer,
                                                    disc_weight=disc_weight,
                                                    filter_p=filter_p,
                                                    target_p=target_p,
                                                    class_bias=class_bias,
                                                    attr_class=attr_class,
                                                    code_0=code_undesired,
                                                    code_1=code_desired,
                                                    multi_code=multi_code,
                                                    )
        else: 
          generated_sequence = model.generate(input_ids=encoded_prompts,
                                                    pad_lens=None,
                                                    max_length=length + length_of_prompt_in_tokens,
                                                    top_k=None,
                                                    top_p=None,
                                                    repetition_penalty=1.2,
                                                    rep_penalty_scale=10,
                                                    eos_token_ids=tokenizer.eos_token_id,
                                                    pad_token_id=0,  # self.tokenizer.eos_token_id,
                                                    do_sample=True,
                                                    temperature = temperature,
                                                    penalize_cond=True,
                                                    gedi_model=gedi_model,
                                                    tokenizer=tokenizer,
                                                    disc_weight=disc_weight,
                                                    filter_p=filter_p,
                                                    target_p=target_p,
                                                    class_bias=class_bias,
                                                    attr_class=attr_class,
                                                    code_0=code_undesired,
                                                    code_1=code_desired,
                                                    multi_code=multi_code,
                                                    gpt3_api_key=gpt3_id,
                                                    )
        text = tokenizer.decode(generated_sequence.tolist()[0], clean_up_tokenization_spaces=True,
                                    skip_special_tokens=True)
  
        text = text[length_of_prompt:]
      else:
        if gpt3_id is None:
          generated_sequence = model.generate(input_ids=encoded_prompts,
                                                    pad_lens=None,
                                                    max_length=length + length_of_prompt_in_tokens,
                                                    top_k=None,
                                                    top_p=None,
                                                    repetition_penalty=1.2,
                                                    rep_penalty_scale=10,
                                                    eos_token_ids=tokenizer.eos_token_id,
                                                    pad_token_id=0,  # self.tokenizer.eos_token_id,
                                                    do_sample=True,
                                                    temperature = temperature, 
                                                    penalize_cond=True,
                                                    gedi_model=None,
                                                    tokenizer=tokenizer,
                                                    disc_weight=disc_weight,
                                                    # filter_p=filter_p,
                                                    # target_p=target_p,
                                                    class_bias=class_bias,
                                                    attr_class=attr_class,
                                                    # code_0=code_undesired,
                                                    # code_1=code_desired,
                                                    # multi_code=multi_code,
                                                    )
          text = tokenizer.decode(generated_sequence.tolist()[0], clean_up_tokenization_spaces=True,
                                    skip_special_tokens=True)
  
          text = text[length_of_prompt:]
          
        else:
          import openai
          openai.api_key = gpt3_id
          completion = openai.Completion()
          response = completion.create(prompt=prompt,
                                   engine="curie",
                                   max_tokens=length,
                                   temperature=temperature,)
          text = response["choices"][0]["text"]
      # if type(extra_args) is dict and 'get_gen_token_count' in extra_args:
      #     return len(generated_sequence.tolist()[0])
  
      
      text = cut_into_sentences(text)
      if len(text) == 0:
          print("Warning! No text generated.")
          return ""
      all_gen_text = text[0]
      return all_gen_text
  
  
  def continuing_generation(prompts, generation_controls, characters, temperature=0.8, gpt3_id=None, disc_weight=30):
    """
    Explanations on controls
    prompts: The prompt to be input. This is a list of sentences. 
    generation_controls: Generation control in the list. If no control is given, -1 is given.
    
    """
    generated = []
  
    prompt_prepend = """###
  Character: Kelly, grandmother
  Start: Kelly found her grandmother's pizza recipe in a shoebox of memories. Kelly reminisced about how much she loved her grandmother's pizza. Kelly decided that she was going to try to make pizza.
  Story after start,: Kelly studied the recipe and gathered everything she needed. Kelly successfully made a pizza from her grandmother's recipe.
  ###
  Character: Leo
  Start: Leo wore a toupee and was anxious about it. Leo decided to go out for a short walk.
  Story after start: It was a very windy day, but he wasn't too concerned. Suddenly, a strong wind came through and took his toupee! His dog leaped and caught it, and he quickly ran home.
  ###
  Character: Jimmy
  Start: Jimmy was a master with his grill. He spent every weekend getting better with his grill. One day he was offered a TV show about grilling on a local channel, Jimmy accepted the job in an instant.
  Story after start: He quit his day job and spent all his time grilling.
  ###
  Character: Mel, Linda
  Start: Mel had a friend, Linda, that Mel didn't know well. Mel let her over my house.
  Story after start: Linda paid rent then asked for some of it back. Linda drinks my juice and eats my food. Linda also makes a huge mess and is very sloppy. Linda got kicked out in two weeks by Mel.
  ###"""
    character_prepend = '\nCharacter:'
    for idx, character in enumerate(characters):
      character_prepend = character_prepend+' '+character
      if idx != len(characters)-1:
        character_prepend = character_prepend + ','
      else:
        character_prepend = character_prepend+'\n'
  
    prompt_start_idx = 0
    for generation_control in generation_controls:
      # print(generation_control)
      prompt_postpend = 'Start: '
      while True:
        for i in range(prompt_start_idx, len(prompts)):
          prompt_postpend = prompt_postpend + prompts[i]
          if i != len(prompts)-1:
            # prompt_postpend = prompt_postpend + ' '
            continue
          else:
            prompt_postpend = prompt_postpend + '\nStory after start:'
        prompt_input = prompt_prepend+character_prepend+prompt_postpend
        prompt_encoded = tokenizer.encode(prompt_input)
        length_of_prompt_in_tokens = len(prompt_encoded)
        if length_of_prompt_in_tokens>1024:
          prompt_start_idx = prompt_start_idx + 1
        else:
          break
      # print(prompt_input)
      gen_sent = generate_one_sentence(prompt_input, generation_control, temperature=temperature, gpt3_id=gpt3_id, disc_weight=disc_weight)
      prompts.append(gen_sent)
      generated.append(gen_sent)
    print(generated)
    return generated