# ---
# 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/"} id="4V30dDc6wn07" outputId="1535b52b-23f8-4392-b28b-31991acb6025"
# !nvidia-smi

# + colab={"base_uri": "https://localhost:8080/"} id="il-IQ1A4UJgP" outputId="4272ac1c-f7ec-4b98-eeea-264522bcc9d7"
# !pip install transformers'==4.2.2'

# + id="_fFHWKDEoIuk"
import torch
from transformers import (
    GPT2Tokenizer, GPT2LMHeadModel, TextDataset, LineByLineTextDataset,
    DataCollatorForLanguageModeling, TrainingArguments, Trainer, get_cosine_schedule_with_warmup
)

import matplotlib.pyplot as plt
# %matplotlib inline

import json
import numpy as np
import seaborn as sns
import json
import nltk


# + colab={"base_uri": "https://localhost:8080/"} id="hh6fjAqjpkrq" outputId="8728dbd0-8248-48c0-f261-7a6108bb2cbb"
# !wget https://s3.amazonaws.com/datasets.huggingface.co/personachat/personachat_self_original.json

# + id="xg0HHhFFRX6k"
bos_token = '<start>'
eos_token = '<end>'

# + id="YD4qBeqOp51m"
with open('/content/personachat_self_original.json') as json_file:
    data = json.load(json_file)

# + id="GJQpDSUejPMh"
chats = []
for i in range(len(data['train'])):
    personality = " ".join(data['train'][i]['personality'])
    chat = data['train'][i]['utterances'][-1]['history']
    final_text = bos_token + personality + '<start_chat>'
    for j in range(len(chat)):
        if j % 2 == 0:
            final_text += '<Human>'
        else:
            final_text += '<AI>'
        final_text += chat[j]
    chats.append(final_text + eos_token)

# + id="jpccGWtrlglm"
chats = " ".join(chats)

# + id="MMr2ug4ylSa3"
train_data_path = 'chats.txt'
with open(train_data_path, 'w') as t:
    t.write(chats)

# + [markdown] id="Y1S4cjsQnLxG"
# # Tokenezation
#

# + colab={"base_uri": "https://localhost:8080/", "height": 98, "referenced_widgets": ["1e4769c47a174dbb80aa31b73d60d6df", "aee2f1e955284382a340e853f05ee468", "f86c9687a0554592996b7c44f77cd447", "<KEY>", "796acd6beaa1474eab01d5e4850802b8", "<KEY>", "86f50e6cdaff4eb885681ba1fa22b888", "e4379d416f1c4a6dab6edf74b02cb1c4", "aab1c2886e2f47c6ad3f1fb51487dc92", "437238a1e32340b5839efbe94a33b5e9", "7bd08e7111ed49488564ee8a8d160520", "7afee4e6d42f41238916e61f35321e1e", "a06b8e6dd8b14d87b1912165339f90db", "cebe23073434472f8b5b1db624b3305b", "<KEY>", "<KEY>", "691e1051e3ee4c1fadf9f7ae20c89f98", "<KEY>", "<KEY>", "1c9ea78649d54f19b167d494c8e4c8f2", "<KEY>", "f69b4b074f254e88aa46de31ce174278"]} id="mkJPniCBUpqQ" outputId="049a8a17-88b3-4895-aa07-1621f6648602"
tokenizer = GPT2Tokenizer.from_pretrained("gpt2", bos_token=bos_token, eos_token=eos_token, pad_token='<|pad|>')

# + id="j8FItmRBWSxS"
block_size = 1024

# + colab={"base_uri": "https://localhost:8080/"} id="jb4j8jmzvOc1" outputId="6cb86235-afc8-42f2-a60f-f37e5280b165"
train_dataset = TextDataset(tokenizer=tokenizer, file_path='/content/chats.txt', block_size=block_size, overwrite_cache=True)

# + colab={"base_uri": "https://localhost:8080/"} id="rjIvTs1DQwf8" outputId="fc48a97a-2da0-483d-d53e-1c7e060d7565"
len(train_dataset)

# + [markdown] id="_O_n3PzynV_Q"
# # Training

# + id="bf3jNlt6r-WK" colab={"base_uri": "https://localhost:8080/", "height": 81, "referenced_widgets": ["8121f206a4c44c3aa1459fd19e7dd3db", "a05b72b4e53843d1916092f9e709db24", "56ee63b4f30445d0a6a3543073232bf8", "ee199dfabc474c829d637be2677e7bef", "3eecd0020ba644dd9932826678f0f340", "fe59676a0ddf448f8caac7ef2137b790", "d9c1a76278834c3ba5eda31245f5013a", "18b9b5a8a00f4f2baca2a7c99d0bac5c", "506d35d3e50c40be831fb61a320048b7", "66d6a8a4d46a43169a663b925de097d6", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "77e9ab36dbe9441ebb3e20cee0d10bf7", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>"]} outputId="31ff7433-757e-4392-baf4-aa35dc3d856f"
model = GPT2LMHeadModel.from_pretrained("gpt2")

# + id="uHOrlUse4wm9"
model.cuda()
None

# + colab={"base_uri": "https://localhost:8080/"} id="BZ84KrfeVTLJ" outputId="a3f0d307-6a81-4c92-a566-834874319427"
model.resize_token_embeddings(len(tokenizer))

# + id="9Nou2PgYEKJr"
learning_rate = 2e-5
batch_size = 1
num_epochs = 4

# + id="-zyz1nIHbsrd"
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

seed = 567
training_args = TrainingArguments(
    output_dir='chats_output',
    overwrite_output_dir=True,
    do_train=True,
    num_train_epochs=num_epochs,
    per_device_train_batch_size=batch_size,
    prediction_loss_only=True,
    logging_steps=100,
    load_best_model_at_end=True,
    #save_steps=300,
    seed=seed,
    learning_rate=learning_rate
)

# + id="76wLGGclb1cQ"
trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    data_collator=data_collator,
    train_dataset=train_dataset
)

train_dataloader = trainer.get_train_dataloader()
num_train_steps = len(train_dataloader) * num_epochs
trainer.create_optimizer_and_scheduler(num_train_steps) 

trainer.lr_scheduler = get_cosine_schedule_with_warmup(
    trainer.optimizer,
    num_warmup_steps=num_train_steps//10,
    num_training_steps=num_train_steps
)

# + [markdown] id="bSdoqrdNDune"
# Заводим!

# + colab={"base_uri": "https://localhost:8080/", "height": 1000} id="Pq5ZPdE1b8c9" outputId="0a1b003f-edda-45af-d9d8-a525b7685db1"
trainer.train() 

# + [markdown] id="kgT-Z-wGFj8O"
# # Save model

# + colab={"base_uri": "https://localhost:8080/"} id="_gFW1bS5Fr3U" outputId="d1a658fc-3ca9-4828-d5ec-dec80c8ca5d0"
output_dir = '/content/chats_output'

model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)


# + [markdown] id="tYcS6UTnE6Gu"
# # Text generation
#

# + id="s8_y37NeJc3u"
def generate(prompt):

    model.eval()

    input_ids = tokenizer.encode(prompt, return_tensors='pt')
    input_ids = input_ids.to('cuda') 
    eos_id = tokenizer.encode(eos_token)[0]

    sample_outputs = model.generate(
        input_ids, 
        max_length= len(input_ids[0]) + 50, 
        min_length=len(input_ids[0]) + 30,
        do_sample=True,
        top_k=50,
        top_p=0.92, 
        temperature=0.95,
        num_beams=1,
        repition_penalty=3,
        eos_token_id=eos_id,
        num_return_sequences=1
        )
    for i, sample_output in enumerate(sample_outputs):
          sample_output = sample_output[len(input_ids[0]):]
          return tokenizer.decode(sample_output, skip_special_tokens=False)

# + id="JZKW0SiNueZ5"


# + id="yPuu-693vFmt"
prompt = '<start>i am a musician and hope to make it big some day . i play the piano and guitar and sing . i also work as a custodian to help pay the bills . my favorite type of music to sing is folk music .<start_chat>'
while True:
    prompt += '<Human>'
    prompt += input('Input: ') + '<AI>'
    generated = generate(prompt)
    if '<Human>' in generated:
        generated = generated[: generated.find('<Human>')]
    print(f'AI: {generated}')
    prompt += generated

# + colab={"base_uri": "https://localhost:8080/"} id="DwF_pKBm88b0" outputId="8d11aa2d-3f1e-4afb-ad8d-cc5c50528739"
from google.colab import drive
drive.mount('/content/drive')

StarcoderdataJupyterScriptsDedupFiltered