<jupyter_start><jupyter_text>Tutorial> API details.<jupyter_code>#hide
from nbdev.showdoc import *
import torch
import wandb

import pytorch_lightning as pl

# from codecarbon import EmissionsTracker
from pathlib import Path
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, TQDMProgressBar
from pytorch_lightning.loggers import WandbLogger
from retrofit.data import RetroDataset
from retrofit.model import RetroFitModelWrapper
from transformers import AutoTokenizer, EncoderDecoderModel
column = "whole_func_string"
encoder_name = "microsoft/codebert-base"
decoder_name = "gpt2"
encoder_tokenizer = AutoTokenizer.from_pretrained(encoder_name)
decoder_tokenizer = AutoTokenizer.from_pretrained(decoder_name)
decoder_tokenizer.pad_token = decoder_tokenizer.eos_token
retro_ds = RetroDataset(
    "code_search_net",
    "flax-sentence-embeddings/st-codesearch-distilroberta-base",
    encoder_tokenizer,
    decoder_tokenizer,
    dataset_config="python",
    column=column,
    batch_size=4,
    k=2,
    n_perc=2
)

model = EncoderDecoderModel.from_encoder_decoder_pretrained(encoder_name, decoder_name)
model.config.decoder_start_token_id = decoder_tokenizer.bos_token_id
model.config.pad_token_id = decoder_tokenizer.pad_token_id
model.config.decoder_max_length = 512
retro_model = RetroFitModelWrapper(
    model,
    weight_decay=0.1,
    lr=5e-4,
    freeze_decoder=True
)
def train_model(model, data_module, num_epochs, output_dir, gradient_accumulation_steps=1):
    """
    Train a model with a given training data loader, validation data loader,
    optimizer, scheduler, loss function, metrics, and callbacks.

    Args:
        model (pl.LightningModule): The model to train.
        data_module (pl.LightningDataModule): The data module to use for training.
        num_epochs (int): The number of epochs to train for.
        output_dir (pathlib.Path): The directory to save the model to.
        name (str): The name of the model.
    Returns:
        best_model_path (str): The path to the best model's checkpoint.
    """
    # pl.seed_everything(115, workers=True)
    wandb_logger = WandbLogger(project="Retrofit")
    # saves a file like: my/path/sample-mnist-epoch=02-val_loss=0.32.ckpt
    checkpoint_callback = ModelCheckpoint(
        monitor="val_loss",
        dirpath=str(output_dir / "checkpoints"),
        filename="retrofit-{epoch:02d}-{val_loss:.2f}",
        save_top_k=5,
        mode="min",
    )
    trainer = pl.Trainer(
        logger=wandb_logger,
        default_root_dir=str(output_dir / "checkpoints"),
        gpus=torch.cuda.device_count(),
        max_epochs=num_epochs,
        limit_train_batches=0.1,
        limit_val_batches=0.1,
        precision=16,
        accumulate_grad_batches=gradient_accumulation_steps,
        callbacks=[
            checkpoint_callback,
            EarlyStopping(monitor="val_loss"),
            TQDMProgressBar(refresh_rate=1),
        ],
    )
    # tracker = EmissionsTracker(output_dir=output_dir.parent.parent, project_name=wandb_logger.name)

    # train the model and track emissions
    # tracker.start()
    trainer.fit(model, data_module)
    # tracker.stop()

    # save the best model to wandb
    best_model_path = checkpoint_callback.best_model_path
    if best_model_path is not None:
        wandb.save(best_model_path)

    # # save the emissions csv file
    # wandb.save(str(output_dir.parent.parent / "emissions.csv"))

    return best_model_path
num_epochs = 5
out_dir = Path("/workspace/retrofit/data/output/")
best_model_path = train_model(
    retro_model,
    retro_ds,
    num_epochs=num_epochs,
    output_dir=out_dir / "model",
    gradient_accumulation_steps=4,
)
# model = RetroFitModel.load_from_checkpoint(best_model_path)
model = RetroFitModelWrapper.load_from_checkpoint(best_model_path)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08595170085692028, 'fraction_numerical': 0.008569202804466372, 'mean_word_length': 4.095238095238095, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '147546', 'license': 'Apache-2.0', 'n_tokens_mistral': 1297, 'n_tokens_neox': 1261, 'n_words': 243, 'path': 'nbs/02_tutorial.ipynb', 'repo_name': 'ncoop57/retrofit', 'size': '3851'}