# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3.8.10 64-bit
#     name: python3
# ---

# + [markdown] id="Ja7zPK-Gvi7Z"
# # Prompt Tuning
#
#

# + id="cE2jNS5UMXKh" colab={"base_uri": "https://localhost:8080/"} outputId="0a6fe469-fc6e-452d-cfba-e1611ce98a2b"
import torch
colab = 'google.colab' in str(get_ipython())

# You need a T4. A K80 will not work.
if colab:
    # !nvidia-smi
    gpu_type = torch.cuda.get_device_name(0)
    if gpu_type != 'Tesla T4':
        raise ValueError("I don't know about this, chief")

# + id="pZDVExyEtR22" colab={"base_uri": "https://localhost:8080/"} outputId="99d30507-f253-4971-ac2e-745ec5f41ba6"
# Setup for Colab only
if colab:
    # !pip install git+https://github.com/finetuneanon/transformers@gpt-neo-localattention3
    # !pip install git+https://github.com/corolla-johnson/mkultra.git#egg=mkultra --log PIP_LOG
    # !pip install gdown
    # !pip install datasets
    # !pip install tqdm

# + colab={"base_uri": "https://localhost:8080/"} id="Pbb7I3GHMXKl" outputId="98169f60-5b94-4df3-ff57-31f20ab3c525"
# If on Colab, mount your Google Drive first!
if colab:
    from google.colab import drive
    drive.mount('/content/drive')

# + id="tWYdCgONMXKm"
# Decide the length of your training blocks in tokens.
# Sizes with headroom for gpt-neo-2.7B-halved:
#  - 700 on a Colab T4 (16GB)
#  - 400 on a Colab K80 (12GB)
#  - 32 on a GTX1080 (8GB)
# If it seems a bit small, don't worry!
# Soft prompts can be moved forward in context for the best effect.
if colab:
    if gpu_type == 'Tesla T4':
        block_size = 700
    else:
        block_size = 400
else:
    block_size = 32

# Name your soft prompt project.
sp_name = 'neuromancer-x-gpt2-optuna'

# Specify the model directory or huggingface name.
if colab:
    model_dir = "/content/drive/MyDrive/models/gpt-neo-2.7B-halved/"
else:
    model_dir = "D:/Git Repos/mkultra/models/gpt-neo-2.7B-halved/" 

# Specify the path to the text file used for training.
if colab:
    text_path = "/content/drive/MyDrive/datasets/neuromancer_reformatted.txt"
else:
    text_path = "datasets/neuromancer_reformatted.txt"

# Specify the project directory bases.
if colab:
    project_dir_root = f"/content/drive/MyDrive/soft_prompts/{sp_name}/"
else:
    project_dir_root = f"soft_prompts/{sp_name}/"

shuffle_seed = 1234567890

eval_percentage = 1.0

# + colab={"base_uri": "https://localhost:8080/", "height": 167, "referenced_widgets": ["88ee0e4be902432d9443b35548198d7d", "<KEY>", "dc9146a630b04dacb084a2f09de921b7", "<KEY>", "<KEY>", "9a0b74d190e84d55965a6906b58dcd2a", "97a649a3ed5c435d9979125d13f7edeb", "<KEY>", "65f505f3e3b84757811463343b3e050e", "7ed9ee26150146f4b0f74f6564da8d00", "<KEY>", "<KEY>", "<KEY>", "e206dc2a2642449c95fd2c2acc4fd00b", "20d7bd96d19743719a949feddd58555c", "<KEY>", "5e7d332212be4a60b2573faa248bd1ac", "<KEY>", "<KEY>", "34b402b0971e4406908a62d4f1aa9aa0", "<KEY>", "<KEY>", "d2dfb62defe14e97b2038e2a284c752c", "<KEY>"]} id="kEVELdEDttvO" outputId="3699af89-924a-412d-9775-47c632cdccd0"
from mkultra.tuning import GPT2PromptTuningLM
model = GPT2PromptTuningLM.from_pretrained(model_dir).half().to("cuda")

# +
import optuna

def objective(trial):
    import os
    from mkultra.trainers import SoftPromptTrainer
    from transformers import Adafactor

    project_dir = os.path.join(project_dir_root, f"{sp_name}-trial-{trial.number}")

    optimizer = Adafactor(
        params=[model.get_soft_params()],
        clip_threshold=trial.suggest_float("clip_threshold", 1.0, 2.0),
        decay_rate=(trial.suggest_float("decay_rate", 0.0, 0.8)-0.8),
        weight_decay=trial.suggest_float("weight_decay", 1e-6, 0.5, log=True),
        beta1=trial.suggest_float("beta1", 0.0, 0.99),
        warmup_init=trial.suggest_categorical("warmup_init", [True, False]))

    n_tokens = trial.suggest_int("n_tokens", 20, 100)

    trainer = SoftPromptTrainer(
        model=model,
        optimizer=optimizer,
        project_dir=project_dir,
        text_path=text_path,
        block_size=1024-n_tokens,
        n_tokens=n_tokens,
        shuffle_seed=shuffle_seed)

    for i in range(5):
        try:
            trainer.shuffle_seed = shuffle_seed+i
            trainer.train()

            trainer.shuffle_seed = shuffle_seed
            eval_loss = trainer.evaluate(eval_percentage=eval_percentage)
            trial.report(eval_loss, i)

            if trial.should_prune():
                raise optuna.TrialPruned()

        except ValueError:
            # Return 10.0 for a NaN
            return 10.0

    return eval_loss


# -

import optuna
study_db = "sqlite:///" + os.path.join(project_dir_root, f"trials.db")
study = optuna.create_study(direction="minimize", storage=study_db, load_if_exists=True)

study.optimize(objective, n_trials=20)

StarcoderdataJupyterScriptsDedupFiltered