# ---
# 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
# ---

# + [markdown] id="MZMbGCw0Qxfg"
# # **Fine-tuning GPT2 for text generation**
# ### This notebook has been taken from [here](https://gist.github.com/GeorgeDittmar/5c57a35332b2b5818e51618af7953351)
# - It contains code to fine tune GPT-2 on two datasets:
#   1. [Tiny Shakespeare](https://huggingface.co/datasets/tiny_shakespeare) - 40,000 lines of Shakespeare from a variety of Shakespeare's plays
#   2. [Bill Sum](https://huggingface.co/datasets/billsum) - Summarization of US Congressional and California state bills.
#

# + [markdown] id="gCoFXKmRS2ru"
# ## Table of contents
# 1. [Imports and Installation section](#Imports)
# 2. [Shakespeare Data preparation](#Shakespeare)
# 3. [GPT-2 Model Fine tuning on Shakespeare Dataset](#FineTuning)
# 4. [Shakespeare Text Generation using GPT-2](#TextGeneration) 
# 5. [Bill Sum Data set preparation](#BillSum)
# 6. [GPT-2 Model Fine tuning on Bill Sum Dataset](#FineTuningBillSum)
# 7. [Bill Sum Text Generation using GPT-2](#BillSumTextGeneration) 

# + [markdown] id="P0DeILhoTAVZ"
# ## Part 1: Imports and Installation section <a name="Imports"></a>

# + id="le99a8fWkpoX"
# import necessary modules
import os
import json
import pandas as pd
from sklearn.model_selection import train_test_split

# + colab={"base_uri": "https://localhost:8080/"} id="B2z-X4SYkwE8" outputId="57846aa6-6d08-49a3-ab22-5b6b6e8eb887"
# mount google drive
from google.colab import drive
drive.mount('/content/drive')

# + colab={"base_uri": "https://localhost:8080/"} id="91rmSAUQVIUP" outputId="edcedec0-a853-42cf-de84-a8a59d1dda90"
#Clone the transformers repo into the notebook
# !git clone https://github.com/huggingface/transformers

# + colab={"base_uri": "https://localhost:8080/"} id="RFzEr4Y1VJKP" outputId="63d4586d-d0cb-4ed1-f21a-0116862ee053"
# Clone should now be in the machine
# !ls

# + [markdown] id="XOTdb4rWv8YN"
# Change directory location to be in the examples folder and then install any requirements

# + colab={"base_uri": "https://localhost:8080/"} id="K2M2Oz9CYB4P" outputId="b3e606e3-c581-4e60-9a8e-51bebd214f0d"
os.chdir("transformers")
os.chdir("./examples/language-modeling")
# !ls

# + colab={"base_uri": "https://localhost:8080/"} id="04rBGxwiYnep" outputId="7f44c125-0afa-456a-9389-69e1ce0f6e07"
# !pip install -r requirements.txt

# + colab={"base_uri": "https://localhost:8080/"} id="iB29mAKjaNIQ" outputId="acfe161c-cfff-4de1-b44f-86256bed26d8"
# !ls

# + colab={"base_uri": "https://localhost:8080/"} id="eo5gRmXaWx0m" outputId="7b53d919-a8cd-44a4-f6bc-d1485af8b19b"
# !pip install pyarrow --upgrade

# + id="y5sdYSpAWY1S"
os.chdir("/content/transformers/examples/")
os.chdir("./language-modeling")

# + colab={"base_uri": "https://localhost:8080/"} id="a6G6WINaYmwx" outputId="bb2dab56-fac3-4d8d-ea0f-180a8251c5d8"
# Need to install latest transformer packages from github so the scripts will run correctly
# !pip install git+git://github.com/huggingface/transformers/

# + [markdown] id="WdWf9TULTUnC"
# ## Part 2: Shakespeare Data Preparation section<a name="Shakespeare"></a>
# - Set up data from a text file in the format <|title|> some data <|endoftext|> and split into training and eval sets.

# + colab={"base_uri": "https://localhost:8080/"} id="uU4ckPTf9T-2" outputId="f1be4580-36a7-4e18-ede1-5f58ac43092c"
with open('/content/drive/MyDrive/bill_sum_shakespeare/shakespeare.txt', 'r') as data:
  dataset = ["<|title|>" + x.strip() for x in data.readlines()]

train, eval = train_test_split(dataset, train_size = 0.9, random_state = 42)
print("Training size:" + str(len(train)))
print("Evaluation size: " + str(len(eval)))

with open('train_tmp.txt', 'w') as file_handle:
  file_handle.write("<|endoftext|>".join(train))

with open('eval_tmp.txt', 'w') as file_handle:
  file_handle.write("<|endoftext|>".join(eval))

# + [markdown] id="1TYi8Oqs6Npo"
# ## Part 3: GPT-2 Model Fine tuning on Shakespeare Dataset<a name="FineTuning"></a>
# ### Fine tuning with **1 epoch**

# + colab={"base_uri": "https://localhost:8080/"} id="yyV_rTL3ZE_H" outputId="9b486ab3-8276-4af7-f3a5-ba3a4ffda41c"
# !python run_clm.py \
# --model_type gpt2-medium \
# --model_name_or_path gpt2-medium \
# --train_file "train_tmp.txt" \
# --do_train \
# --validation_file "eval_tmp.txt" \
# --do_eval \
# --per_gpu_train_batch_size 1 \
# --save_steps -1 \
# --num_train_epochs 1 \
# --fp16 \
# --learning_rate 5e-5 \
# --output_dir="Output dir"

# + [markdown] id="4mctfrJ65FOY"
# ### Fine tuning with **3 epochs**

# + colab={"base_uri": "https://localhost:8080/"} id="uTVGqlyd5ITu" outputId="0d6ead37-066d-4706-81d4-ab3f2fe8b253"
# !python run_clm.py \
# --model_type gpt2-medium \
# --model_name_or_path gpt2-medium \
# --train_file "train_tmp.txt" \
# --do_train \
# --validation_file "eval_tmp.txt" \
# --do_eval \
# --per_gpu_train_batch_size 1 \
# --save_steps -1 \
# --num_train_epochs 3 \
# --fp16 \
# --learning_rate 5e-5 \
# --lr_scheduler_type constant \
# --evaluation_strategy epoch \
# --output_dir="Output dir" \
# --overwrite_output_dir True

# + [markdown] id="QPFqA6NnZ_fM"
# ### Fine tuning with **5 epochs**
#

# + colab={"base_uri": "https://localhost:8080/"} id="uBa9HkqmaEsB" outputId="8ec8385d-96f8-48c8-b932-109a9d78b10c"
# !python run_clm.py \
# --model_type gpt2-medium \
# --model_name_or_path gpt2-medium \
# --train_file "train_tmp.txt" \
# --do_train \
# --validation_file "eval_tmp.txt" \
# --do_eval \
# --per_gpu_train_batch_size 1 \
# --save_steps -1 \
# --num_train_epochs 5 \
# --fp16 \
# --learning_rate 5e-5 \
# --lr_scheduler_type constant \
# --evaluation_strategy epoch \
# --output_dir="Output dir" \
# --overwrite_output_dir True

# + [markdown] id="GbgV2Fe6hJ6n"
# ### Fine tuning with **7 epochs**

# + colab={"base_uri": "https://localhost:8080/"} id="Q4W6zZ8thXS1" outputId="801fc9a2-cbc8-4be6-e96a-1c07eb620334"
# !python run_clm.py \
# --model_type gpt2-medium \
# --model_name_or_path gpt2-medium \
# --train_file "train_tmp.txt" \
# --do_train \
# --validation_file "eval_tmp.txt" \
# --do_eval \
# --per_gpu_train_batch_size 1 \
# --save_steps -1 \
# --num_train_epochs 7 \
# --fp16 \
# --learning_rate 5e-5 \
# --lr_scheduler_type constant \
# --evaluation_strategy epoch \
# --output_dir="Output dir" \
# --overwrite_output_dir True

# + [markdown] id="cyQIU3eXqzha"
# ### Fine tuning with **Learning rate = 4e-5** & number of **epochs = 5**

# + colab={"base_uri": "https://localhost:8080/"} id="dIzrmprOq9mJ" outputId="93244162-e77d-407a-f8c9-b0163d852606"
# !python run_clm.py \
# --model_type gpt2-medium \
# --model_name_or_path gpt2-medium \
# --train_file "train_tmp.txt" \
# --do_train \
# --validation_file "eval_tmp.txt" \
# --do_eval \
# --per_gpu_train_batch_size 1 \
# --save_steps -1 \
# --num_train_epochs 5 \
# --fp16 \
# --learning_rate 4e-5 \
# --lr_scheduler_type constant \
# --evaluation_strategy epoch \
# --output_dir="Output dir" \
# --overwrite_output_dir True

# + [markdown] id="f9asJ-TnwZN3"
# #### Fine tuning with **Learning rate = 3e-5** & number of **epochs = 5**

# + colab={"base_uri": "https://localhost:8080/"} id="twlSvF4sweOF" outputId="42581888-fe40-407e-acf1-2893618856bf"
# !python run_clm.py \
# --model_type gpt2-medium \
# --model_name_or_path gpt2-medium \
# --train_file "train_tmp.txt" \
# --do_train \
# --validation_file "eval_tmp.txt" \
# --do_eval \
# --per_gpu_train_batch_size 1 \
# --save_steps -1 \
# --num_train_epochs 5 \
# --fp16 \
# --learning_rate 3e-5 \
# --lr_scheduler_type constant \
# --evaluation_strategy epoch \
# --output_dir="Output dir" \
# --overwrite_output_dir True

# + [markdown] id="YBIuyEcmwiZF"
# #### Fine tuning with **Learning rate = 2e-5** & number of **epochs = 5**

# + colab={"base_uri": "https://localhost:8080/"} id="r413bD3Jwmdg" outputId="7ce43437-2c1d-4a01-ee60-26c07d000ea4"
# !python run_clm.py \
# --model_type gpt2-medium \
# --model_name_or_path gpt2-medium \
# --train_file "train_tmp.txt" \
# --do_train \
# --validation_file "eval_tmp.txt" \
# --do_eval \
# --per_gpu_train_batch_size 1 \
# --save_steps -1 \
# --num_train_epochs 5 \
# --fp16 \
# --learning_rate 2e-5 \
# --lr_scheduler_type constant \
# --evaluation_strategy epoch \
# --output_dir="Output dir" \
# --overwrite_output_dir True

# + colab={"base_uri": "https://localhost:8080/", "height": 218, "referenced_widgets": ["07ac858352fc4e07a60eaa84310cc5a1", "<KEY>", "dac659c527164f37b209e7aef340fb5b", "c45243609ad244f3aeceb11d3e213bbc", "56a388a1507845bc9537b3e60d9e7477", "<KEY>", "684f3b3ad6284b80ac719c34e566a4b1", "<KEY>", "<KEY>", "7b9dff47948a476183fc396e06b97b96", "e06d435e204d46d2a900142f78a9772a", "<KEY>", "be160c60e7794714ad2bdd1a08c5b1a5", "ed44fe544eab46778cbe6e0d02acc73b", "e40a9cf9e1af47f29b1c97ec750b0469", "<KEY>"]} id="iMDmm7TEVpb4" outputId="dcc1534d-16b0-4e28-efab-47786ba83a31"
from transformers import TFGPT2LMHeadModel
from transformers import GPT2Tokenizer

model = TFGPT2LMHeadModel.from_pretrained("Output dir", from_pt = True)
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

# + [markdown] id="NEpvmAoWB2qh"
# ## Part 4: Shakespeare Text Generation using GPT-2 <a name="TextGeneration"></a>

# + colab={"base_uri": "https://localhost:8080/"} id="bGXQxXhqV2Vz" outputId="dfa225da-a79a-4292-8587-85a1865064fa"
# set the initial text to start the process of text generation
input_ids = tokenizer.encode("He that will give good words to thee will", return_tensors = 'tf')

# print the tensor ids
input_ids[0]

# + colab={"base_uri": "https://localhost:8080/"} id="KCyei38MWUlr" outputId="78c5ead8-63a9-4890-c94a-f069db77d165"
generated_text_samples = model.generate(
    input_ids, 
    max_length = 128,  
    num_return_sequences = 5,
    no_repeat_ngram_size = 2,
    repetition_penalty = 1.5,
    top_p = 0.92,
    temperature = 0.85,
    do_sample = True,
    top_k = 125,
    early_stopping = True
)

# + colab={"base_uri": "https://localhost:8080/"} id="4ChFrv99WiiC" outputId="c63036c1-e65e-4ad8-c37a-1ac26f3d7653"
# Print output for each sequence generated above
for i, token in enumerate(generated_text_samples):
  print("{}: {}".format(i, tokenizer.decode(token, skip_special_tokens = True)))
  print()

# + [markdown] id="Esan38nSBrNc"
# ## Part 5: Bill Sum Data set preparation<a name="BillSum"></a>

# + colab={"base_uri": "https://localhost:8080/", "height": 1000} id="_LxGI3EzBqS-" outputId="2ffec9fe-9d59-446e-fbc7-c507a23c1a30"
# read the training data into pandas dataframe
data_path = '/content/drive/MyDrive/bill_sum_shakespeare/us_train_data_final_OFFICIAL.jsonl'
data = pd.read_json(data_path, lines = True)

# display settings
pd.set_option('display.max_colwidth', None)

# display few rows of the dataframe
data.head(2)

# + colab={"base_uri": "https://localhost:8080/"} id="XyEf-zaSofdL" outputId="3bd545d4-7e9b-48c1-a76a-f13b5b6fde4a"
# check the shape of the data frame
data.shape

# + colab={"base_uri": "https://localhost:8080/"} id="YjS-kNxtpSZj" outputId="065fd0f8-c825-4955-c68f-4068ee078e66"
# extract the text data
text = data['summary'].values.tolist()

# split data into train and validation set
train, eval = train_test_split(text, train_size = 0.9, random_state = 42)
print("Training size:" + str(len(train)))
print("Evaluation size: " + str(len(eval)))

with open('train_tmp.txt', 'w') as file_handle:
  file_handle.write("<|endoftext|>".join(train))

with open('eval_tmp.txt', 'w') as file_handle:
  file_handle.write("<|endoftext|>".join(eval))

# + [markdown] id="AjDqyrOKtHt3"
# ## Part 6: GPT-2 Model Fine tuning on Bill Sum Dataset<a name="FineTuningBillSum"></a>
# ### Fine tuning with **1 epoch**

# + colab={"base_uri": "https://localhost:8080/"} id="5s6-a617tACP" outputId="d0761fde-0130-4e3a-ccdd-7875b990d2ad"
# !python run_clm.py \
# --model_type gpt2-medium \
# --model_name_or_path gpt2-medium \
# --train_file "train_tmp.txt" \
# --do_train \
# --validation_file "eval_tmp.txt" \
# --do_eval \
# --per_gpu_train_batch_size 1 \
# --save_steps -1 \
# --num_train_epochs 1 \
# --fp16 \
# --learning_rate 5e-5 \
# --lr_scheduler_type constant \
# --evaluation_strategy epoch \
# --output_dir="Output dir" \
# --overwrite_output_dir True

# + [markdown] id="Yc850DIAs3az"
# ### Fine tuning with **3 epochs**

# + colab={"base_uri": "https://localhost:8080/"} id="WAUB7Bsgs850" outputId="50c94e30-f998-448f-a1cb-e93357a0b26f"
# !python run_clm.py \
# --model_type gpt2-medium \
# --model_name_or_path gpt2-medium \
# --train_file "train_tmp.txt" \
# --do_train \
# --validation_file "eval_tmp.txt" \
# --do_eval \
# --per_gpu_train_batch_size 1 \
# --save_steps -1 \
# --num_train_epochs 3 \
# --fp16 \
# --learning_rate 5e-5 \
# --lr_scheduler_type constant \
# --evaluation_strategy epoch \
# --output_dir="Output dir" \
# --overwrite_output_dir True

# + [markdown] id="e-aNrkxFLEmF"
# ### Fine tuning with **5 epochs**

# + colab={"base_uri": "https://localhost:8080/"} id="PlVMTR4yLBUO" outputId="79983f0b-dabf-4755-82dd-ddac1d694ff7"
# !python run_clm.py \
# --model_type gpt2-medium \
# --model_name_or_path gpt2-medium \
# --train_file "train_tmp.txt" \
# --do_train \
# --validation_file "eval_tmp.txt" \
# --do_eval \
# --per_gpu_train_batch_size 1 \
# --save_steps -1 \
# --num_train_epochs 5 \
# --fp16 \
# --learning_rate 5e-5 \
# --lr_scheduler_type constant \
# --evaluation_strategy epoch \
# --output_dir="Output dir" \
# --overwrite_output_dir True

# + [markdown] id="3CpzI5mU1jPl"
# ## Part 7: Bill Sum Text Generation using GPT-2 <a name="BillSumTextGeneration"></a>
#

# + colab={"base_uri": "https://localhost:8080/"} id="6aoBsCE0W0ch" outputId="1fa9f987-22da-41d2-e216-81f3dfcfcc34"
model = TFGPT2LMHeadModel.from_pretrained("Output dir", from_pt = True)

# set the initial text to start the process of text generation
bill_sum_input_ids = tokenizer.encode("This Act may be cited as the ``Farm to School Improvements Act", return_tensors = 'tf')

# print the tensor ids
bill_sum_input_ids[0]

# + [markdown] id="BZ3YNTNlBsh8"
# ### Generate Text

# + colab={"base_uri": "https://localhost:8080/"} id="NbzHNvvaAPns" outputId="d3cb0c22-b425-4a7d-ef26-c3140c93925a"
bill_sum_generated_text_samples = model.generate(
    bill_sum_input_ids, 
    max_length = 128,  
    num_return_sequences = 5,
    no_repeat_ngram_size = 2,
    repetition_penalty = 1.5,
    top_p = 0.92,
    temperature = 0.85,
    do_sample = True,
    top_k = 125,
    early_stopping = True
)

# + colab={"base_uri": "https://localhost:8080/"} id="XPdteSR_B3w1" outputId="e662b947-4962-47ae-86e7-8ddfb967801d"
# Print output for each sequence generated above
for i, token in enumerate(bill_sum_generated_text_samples):
  print("{}: {}".format(i, tokenizer.decode(token, skip_special_tokens = True)))
  print()

StarcoderdataJupyterScriptsDedupFiltered