# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + colab={"base_uri": "https://localhost:8080/", "height": 185} colab_type="code" id="--jFBFXI64sl" outputId="84e30ccb-a173-402f-e030-0f7920a9c881"
# !pip install -q git+https://github.com/huggingface/transformers.git
# !pip install -q tensorflow==2.1

# + colab={"base_uri": "https://localhost:8080/", "height": 50} colab_type="code" id="lVQzctdr6_3n" outputId="b23ea2c7-42e5-4e40-bd4d-6b312c338b6e"
import torch
import pandas as pd
import nltk
nltk.download('punkt')
from nltk import tokenize
import logging
from tqdm import tqdm
from tqdm import trange
from transformers import (
    CTRLLMHeadModel,
    CTRLTokenizer,
    GPT2LMHeadModel,
    GPT2Tokenizer,
    OpenAIGPTLMHeadModel,
    OpenAIGPTTokenizer,
    TransfoXLLMHeadModel,
    TransfoXLTokenizer,
    XLMTokenizer,
    XLMWithLMHeadModel,
    XLNetLMHeadModel,
    XLNetTokenizer,
)

def prepare_ctrl_input(temperature, _, tokenizer, prompt_text):
    if temperature > 0.7:
        logger.info("CTRL typically works better with lower temperatures (and lower top_k).")

    encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False)
    if not any(encoded_prompt[0] == x for x in tokenizer.control_codes.values()):
        logger.info("WARNING! You are not starting your generation from a control code so you won't get good results")
    return prompt_text


def prepare_xlm_input(xlm_language, model, tokenizer, prompt_text):
    # kwargs = {"language": None, "mask_token_id": None}

    # Set the language
    use_lang_emb = hasattr(model.config, "use_lang_emb") and model.config.use_lang_emb
    if hasattr(model.config, "lang2id") and use_lang_emb:
        available_languages = model.config.lang2id.keys()
        if xlm_language in available_languages:
            language = xlm_language
        else:
            language = None
            while language not in available_languages:
                language = input("Using XLM. Select language in " + str(list(available_languages)) + " >>> ")

        model.config.lang_id = model.config.lang2id[language]
        # kwargs["language"] = tokenizer.lang2id[language]

    # TODO fix mask_token_id setup when configurations will be synchronized between models and tokenizers
    # XLM masked-language modeling (MLM) models need masked token
    # is_xlm_mlm = "mlm" in args.model_name_or_path
    # if is_xlm_mlm:
    #     kwargs["mask_token_id"] = tokenizer.mask_token_id

    return prompt_text


def prepare_xlnet_input(padding_text, _, tokenizer, prompt_text):
    prompt_text = (padding_text if padding_text else PADDING_TEXT) + prompt_text
    return prompt_text


def prepare_transfoxl_input(padding_text, _, tokenizer, prompt_text):
    prompt_text = (padding_text if padding_text else PADDING_TEXT) + prompt_text
    return prompt_text


PREPROCESSING_FUNCTIONS = {
    "ctrl": prepare_ctrl_input,
    "xlm": prepare_xlm_input,
    "xlnet": prepare_xlnet_input,
    "transfo-xl": prepare_transfoxl_input,
}

MAX_LENGTH = int(10000)

def adjust_length_to_model(length, max_sequence_length):
    if length < 0 and max_sequence_length > 0:
        length = max_sequence_length
    elif 0 < max_sequence_length < length:
        length = max_sequence_length  # No generation bigger than model size
    elif length < 0:
        length = MAX_LENGTH  # avoid infinite loop
    return length


device = torch.device("cuda" if torch.cuda.is_available()  else "cpu")


# + colab={} colab_type="code" id="0oKs1wk8QpP7"
def initialize_model(model_type, length):

  MODEL_CLASSES = {
      "gpt2": (GPT2LMHeadModel, GPT2Tokenizer),
      "ctrl": (CTRLLMHeadModel, CTRLTokenizer),
      "openai-gpt": (OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
      "xlnet": (XLNetLMHeadModel, XLNetTokenizer),
      "transfo-xl": (TransfoXLLMHeadModel, TransfoXLTokenizer),
      "xlm": (XLMWithLMHeadModel, XLMTokenizer),
  }
  try:
      model_class, tokenizer_class = MODEL_CLASSES[model_type]
  except KeyError:
      raise KeyError("the model {} you specified is not supported.")
  tokenizer = tokenizer_class.from_pretrained(model_type)
  model = model_class.from_pretrained(model_type)
  model.to(device)
  length = adjust_length_to_model(length, max_sequence_length=model.config.max_position_embeddings)
  return model, tokenizer, length


# + colab={"base_uri": "https://localhost:8080/", "height": 212, "referenced_widgets": ["6b80adae7498437790059c6f05d65a0d", "d20830f40f354d80b3ba5dce85fc6a9e", "dd4c9167792b4c7eb0b820abd26269e9", "9f1b97869062457eb7bd8671127a5514", "<KEY>", "526e140d6e7c438bab1386f8c461efed", "6cc4a19887444edb932d67f08a43754e", "<KEY>", "8955e7123c524a2c9a0f84e38dfe1327", "79e581b72a924617922892bab1b7db00", "<KEY>", "<KEY>", "<KEY>", "48966cc8360845e0b06ede125339123b", "<KEY>", "b224ab9cede244bdb66a1ca13cf36768", "<KEY>", "edaf86a6fbc9483dbdc3e44e414c42e7", "<KEY>", "<KEY>", "9bd8d9ed81634d5fa77869d189b7437c", "b23132cadc98454797a781e8c9a669e7", "3462eda4bdea4b8396fa574a786cf31a", "203bfe73ef1a4e3f9ec77a825c4b7962", "a13600dba76a417da31ca1f7efdd08d6", "13334137efa147da840a0f857fc8321c", "2de98aa0244948e2813f4105ab9ac7f7", "6263c4be54d843ba9ede0208c3ceb3cd", "e39a7484b27346879ee32a3296c62128", "<KEY>", "781a434546bb402a8e9d812ae87e2c38", "3562fdbf3e404b13b67a0068c7785a70"]} colab_type="code" id="tp2JuYeEBgxb" outputId="29606dca-4d2e-485a-a311-fbfaeea49028"
model_type = "ctrl" #PUT NAME OF NEEDED MODEL
length = 200 #MODEL LENGTH
prompt_text = "Links Java vs Python"
repetition_penalty = 1.2
temperature = 0.2
stop_token = None
num_return_sequences = 5
# Initialize the model and tokenizer
model, tokenizer, length =  initialize_model(model_type, length)


# + colab={"base_uri": "https://localhost:8080/", "height": 374} colab_type="code" id="-28paeXrU8jD" outputId="ea83197d-a0f6-476f-e896-3df2f4ed3bb5"
def generate_text_from_condition(model, tokenizer, length, prompt_text, repetition_penalty, temperature, num_return_sequences, model_name, stop_token = None):
   # Initialize the model and tokenizer
    # prompt_text = args.prompt if args.prompt else input("Model prompt >>> ")

    # Different models need different input formatting and/or extra arguments
  requires_preprocessing = model_type in PREPROCESSING_FUNCTIONS.keys()
  if requires_preprocessing:
      prepare_input = PREPROCESSING_FUNCTIONS.get(model_type)
      preprocessed_prompt_text = prepare_input(temperature, model, tokenizer, prompt_text)
      encoded_prompt = tokenizer.encode(
          preprocessed_prompt_text, add_special_tokens=False, return_tensors="pt", add_space_before_punct_symbol=True
      )
  else:
      encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
  encoded_prompt = encoded_prompt.to(device)

  if encoded_prompt.size()[-1] == 0:
      input_ids = None
  else:
      input_ids = encoded_prompt

  output_sequences = model.generate(
      input_ids=input_ids,
      max_length=length + len(encoded_prompt[0]),
      temperature=temperature,
      # top_k=args.k,
      # top_p=args.p,
      repetition_penalty=repetition_penalty,
      do_sample=True,
      num_return_sequences=num_return_sequences,
  )

  # Remove the batch dimension when returning multiple sequences
  if len(output_sequences.shape) > 2:
      output_sequences.squeeze_()

  generated_sequences = []

  for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
      generated_sequence = generated_sequence.tolist()

      # Decode text
      text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)

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

      # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
      total_sequence = (
          prompt_text + text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)) :]
      )
      if model_name == 'ctrl':
        total_sequence = total_sequence.replace(prompt_text + " \n \n ", "")
        total_sequence = " ".join(tokenize.sent_tokenize(total_sequence)[:-1])
      else:
        total_sequence = total_sequence.replace(prompt_text + "\n", "")
        total_sequence = " ".join(tokenize.sent_tokenize(total_sequence)[:-1])
      generated_sequences.append(total_sequence)
  return generated_sequences

gen = generate_text_from_condition(model, tokenizer, length, "Links What is better: Java or Python?", repetition_penalty, temperature, num_return_sequences, 'gpt-2', stop_token="<|endoftext|>")
gen

# + colab={} colab_type="code" id="vaZDdt-6ZVd4"
df = pd.read_csv("data\100_short.csv")
df = df[["object_a", 'object_b']]
df = df.reindex(df.columns.tolist() + ['CTRL_gen_0','CTRL_gen_1', "CTRL_gen_2", "CTRL_gen_3", "CTRL_gen_4", 'GPT_2_gen_0','GPT_2_gen_1', "GPT_2_gen_2", "GPT_2_gen_3", "GPT_2_gen_4"], axis=1)

# + colab={"base_uri": "https://localhost:8080/", "height": 573} colab_type="code" id="JDP8n6N0dM6P" outputId="e110b0b4-c420-4a1c-8cbd-ae03af05b01d"
question = "Links What is better: {} or {}?"
for i in trange(df.shape[0]):
  obj_a, obj_b = df.iloc[i, 0:2].tolist()
  prompt = question.format(obj_a, obj_b)
  generated_sequences = generate_text_from_condition(model, tokenizer, length, prompt, repetition_penalty, temperature, num_return_sequences, stop_token=None)
  for j in range(num_return_sequences):
    row_ind = df.index[i]
    df.loc[row_ind, 'CTRL_gen_' + str(j)] = generated_sequences[j]

df

# + colab={} colab_type="code" id="FmjmnRNpHENF"
model_type = "gpt2" #PUT NAME OF NEEDED MODEL
length = 200 #MODEL LENGTH
prompt_text = "Links Java vs Python"
repetition_penalty = 1.2
temperature = 0.2
stop_token = None
num_return_sequences = 5
# Initialize the model and tokenizer
model, tokenizer, length =  initialize_model(model_type, length)

# + colab={"base_uri": "https://localhost:8080/", "height": 342, "referenced_widgets": ["fd9bf7df242a4f86a723d6499ff8c050", "e9e438571a314ccfb834bff82c57987c", "8922308a5acc4cb288a1ba029d5ebadd", "64f43be4bd0b47409eafedae757c32ff", "8d7a93fdff0e4d528b048e7660962941", "<KEY>", "6c958e8b6026434a87cf2e9e904eaf43", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "04912ba2e1b64277b6610e10add9245b", "459d5c6d3f4d4178804daf05adfa8e53", "5c60bac1aa764a7db7479f2adb08e82f", "3222961cb81a436da459b949c7e79629", "eda87cb1f59a4cffa4e263dcedf8799e", "4ed7de85331b49bba20cce8b1a76cab8", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "97573bad9f4c4ea1861f19af2c217820", "86ffdd924bb148bc882726ea43a9edef", "e0c569e8545d46c7aabe04d687e97b47", "<KEY>", "<KEY>", "<KEY>", "1e8a303dd5e8422e9d10b4e90072bb71", "1e8fa00d5a27471c8bc8037091478130", "1f2e94c93ee1481ca3836123e0a5fd34"]} colab_type="code" id="K1tRFCduHN7h" outputId="794efc1f-2578-4943-dcd4-5e33ba6237b2"
model, tokenizer, length =  initialize_model(model_type, length)
gen = generate_text_from_condition(model, tokenizer, length, "What is better: Java or Python?", repetition_penalty, temperature, num_return_sequences, stop_token=None)
gen

# + colab={"base_uri": "https://localhost:8080/", "height": 1000} colab_type="code" id="-oUZ3XmBGlVI" outputId="53646d36-274f-4bf9-962c-e49078d9ca67"
question = "What is better: {} or {}?"
for i in trange(df.shape[0]):
  obj_a, obj_b = df.iloc[i, 0:2].tolist()
  prompt = question.format(obj_a, obj_b)
  generated_sequences = generate_text_from_condition(model, tokenizer, length, prompt, repetition_penalty, temperature, num_return_sequences, model_name='gpt-2', stop_token="<|endoftext|>")
  for j in range(num_return_sequences):
    row_ind = df.index[i]
    df.loc[row_ind, 'GPT_2_gen_' + str(j)] = generated_sequences[j]

df

# + colab={} colab_type="code" id="_S8-XAPZ345P"
df.to_csv("ctrl_gpt.csv")

# + colab={"base_uri": "https://localhost:8080/", "height": 203} colab_type="code" id="mvHAmXcOIrzi" outputId="e85ef701-cd8a-495f-ff33-24e7c79c4840"
d=[]
while(1):
  d.append('1')

StarcoderdataJupyterScriptsDedupFiltered