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

# # Huggingface SageMaker-SDK - T5 Summarization example

# 1. [Introduction](#Introduction)  
# 2. [Development Environment and Permissions](#Development-Environment-and-Permissions)
#     1. [Installation](#Installation)  
#     2. [Permissions](#Permissions)
#     3. [Uploading data to sagemaker_session_bucket](#Uploading-data-to-sagemaker_session_bucket)  
# 3. [Fine-tuning & starting Sagemaker Training Job](#Fine-tuning-\&-starting-Sagemaker-Training-Job)  
#     1. [Creating an Estimator and start a training job](#Creating-an-Estimator-and-start-a-training-job)  
#     2. [Download fine-tuned model from s3](#Download-fine-tuned-model-from-s3)
#     3. [Summarization on Local](#Summarization-on-Local)  

# # Introduction
#
# このnotebookはHuggingFaceの[run_summarization.py](https://github.com/huggingface/transformers/blob/master/examples/pytorch/summarization/run_summarization.py)を日本語データで実行します。    
#
# **51行目の`check_min_version("4.10.0.dev0")`をコメントアウトしたことと、494-509行目の`postprocess_text`を日本語で評価できるように変更しましたが、それ以外はSageMakerで実行するために変更を加えた部分はありません**
#
# 事前訓練モデルには[sonoisa/t5-base-japanese](https://huggingface.co/sonoisa/t5-base-japanese)を使用し、データは[wikiHow日本語要約データセット](https://github.com/Katsumata420/wikihow_japanese)を使用します。    
#
# このデモでは、AmazonSageMakerのHuggingFace Estimatorを使用してSageMakerのトレーニングジョブを実行します。    
#
# _**NOTE: このデモは、SagemakerNotebookインスタンスで動作検証しています**_    

# # Development Environment and Permissions

# ## Installation
#
# このNotebookはSageMakerの`conda_pytorch_p36`カーネルを利用しています。    
#
# **_Note: このnotebook上で推論テストを行う場合、（バージョンが古い場合は）pytorchのバージョンアップが必要になります。_**

# !pip install --upgrade pip
# !pip install --upgrade torch
# !pip install "sagemaker>=2.48.1" "transformers==4.9.2" "datasets[s3]>=1.8.0" --upgrade
# !pip install sentencepiece

# ## Permissions
#
# ローカル環境でSagemakerを使用する場合はSagemakerに必要な権限を持つIAMロールにアクセスする必要があります。[こちら](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html)を参照してください

# +
import sagemaker

sess = sagemaker.Session()
# sagemaker session bucket -> used for uploading data, models and logs
# sagemaker will automatically create this bucket if it not exists
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    # set to default bucket if a bucket name is not given
    sagemaker_session_bucket = sess.default_bucket()

role = sagemaker.get_execution_role()
sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

print(f"sagemaker role arn: {role}")
print(f"sagemaker bucket: {sess.default_bucket()}")
print(f"sagemaker session region: {sess.boto_region_name}")
# -

# # データの準備
#
# 事前に`create_wikihow_dataset.ipynb`を実行してwikiHow日本語要約データセットを用意してください。

import pandas as pd
from tqdm import tqdm

train = pd.read_json('./wikihow_japanese/data/output/train.jsonl', orient='records', lines=True)
train

dev = pd.read_json('./wikihow_japanese/data/output/dev.jsonl', orient='records', lines=True)
dev

test = pd.read_json('./wikihow_japanese/data/output/test.jsonl', orient='records', lines=True)
test

# [run_summarization.py](https://github.com/huggingface/transformers/blob/master/examples/pytorch/summarization/run_summarization.py)の実行にあたり、データを少し修正します。    
# データのフォーマットは`csv`もしくは`json`で以下のように「原文（Text）」と「要約（summary）」を含む必要があります。    
#
# ```
# text,summary
# "I'm sitting here in a boring room. It's just another rainy Sunday afternoon. I'm wasting my time I got nothing to do. I'm hanging around I'm waiting for you. But nothing ever happens. And I wonder","I'm sitting in a room where I'm waiting for something to happen"
# "I see trees so green, red roses too. I see them bloom for me and you. And I think to myself what a wonderful world. I see skies so blue and clouds so white. The bright blessed day, the dark sacred night. And I think to myself what a wonderful world.","I'm a gardener and I'm a big fan of flowers."
# "Christmas time is here. Happiness and cheer. Fun for all that children call. Their favorite time of the year. Snowflakes in the air. Carols everywhere. Olden times and ancient rhymes. Of love and dreams to share","It's that time of year again."
# ```
#
# ```
# {"text": "I'm sitting here in a boring room. It's just another rainy Sunday afternoon. I'm wasting my time I got nothing to do. I'm hanging around I'm waiting for you. But nothing ever happens. And I wonder", "summary": "I'm sitting in a room where I'm waiting for something to happen"}
# {"text": "I see trees so green, red roses too. I see them bloom for me and you. And I think to myself what a wonderful world. I see skies so blue and clouds so white. The bright blessed day, the dark sacred night. And I think to myself what a wonderful world.", "summary": "I'm a gardener and I'm a big fan of flowers."}
# {"text": "Christmas time is here. Happiness and cheer. Fun for all that children call. Their favorite time of the year. Snowflakes in the air. Carols everywhere. Olden times and ancient rhymes. Of love and dreams to share", "summary": "It's that time of year again."}
# ```
#
# 詳細は[こちら](https://github.com/huggingface/transformers/tree/master/examples/pytorch/summarization)にあります。    
# ここでは`jsonl`のファイルを`csv`に変換します。

train.to_csv('train.csv', index=False)
dev.to_csv('dev.csv', index=False)
test.to_csv('test.csv', index=False)

# ## Uploading data to `sagemaker_session_bucket`
#
# S3へデータをアップロードします。

# +
s3_prefix = 'samples/datasets/wikihow_csv'

input_train = sess.upload_data(
    path='train.csv', 
    key_prefix=f'{s3_prefix}/train'
)

input_validation = sess.upload_data(
    path='dev.csv', 
    key_prefix=f'{s3_prefix}/valid'
)

input_test = sess.upload_data(
    path='test.csv', 
    key_prefix=f'{s3_prefix}/test'
)

# +
# データのUpload path

print(input_train)
print(input_validation)
print(input_test)
# -

# # Fine-tuning & starting Sagemaker Training Job
#
# `HuggingFace`のトレーニングジョブを作成するためには`HuggingFace` Estimatorが必要になります。    
# Estimatorは、エンドツーエンドのAmazonSageMakerトレーニングおよびデプロイタスクを処理します。 Estimatorで、どのFine-tuningスクリプトを`entry_point`として使用するか、どの`instance_type`を使用するか、どの`hyperparameters`を渡すかなどを定義します。
#
#
# ```python
# huggingface_estimator = HuggingFace(
#     entry_point='train.py',
#     source_dir='./scripts',
#     base_job_name='huggingface-sdk-extension',
#     instance_type='ml.p3.2xlarge',
#     instance_count=1,
#     transformers_version='4.4',
#     pytorch_version='1.6',
#     py_version='py36',
#     role=role,
#     hyperparameters={
#         'epochs': 1,
#         'train_batch_size': 32,
#         'model_name':'distilbert-base-uncased'
#     }
# )
# ```
#
# SageMakerトレーニングジョブを作成すると、SageMakerは`huggingface`コンテナを実行するために必要なec2インスタンスの起動と管理を行います。    
# Fine-tuningスクリプト`train.py`をアップロードし、`sagemaker_session_bucket`からコンテナ内の`/opt/ml/input/data`にデータをダウンロードして、トレーニングジョブを実行します。
#
#
# ```python
# /opt/conda/bin/python train.py --epochs 1 --model_name distilbert-base-uncased --train_batch_size 32
# ```
#
# `HuggingFace estimator`で定義した`hyperparameters`は、名前付き引数として渡されます。
#
# またSagemakerは、次のようなさまざまな環境変数を通じて、トレーニング環境に関する有用なプロパティを提供しています。
#
# * `SM_MODEL_DIR`：トレーニングジョブがモデルアーティファクトを書き込むパスを表す文字列。トレーニング後、このディレクトリのアーティファクトはモデルホスティングのためにS3にアップロードされます。
#
# * `SM_NUM_GPUS`：ホストで使用可能なGPUの数を表す整数。
#
# * `SM_CHANNEL_XXXX`：指定されたチャネルの入力データを含むディレクトリへのパスを表す文字列。たとえば、HuggingFace estimatorのfit呼び出しで`train`と`test`という名前の2つの入力チャネルを指定すると、環境変数`SM_CHANNEL_TRAIN`と`SM_CHANNEL_TEST`が設定されます。
#
# このトレーニングジョブをローカル環境で実行するには、`instance_type='local'`、GPUの場合は`instance_type='local_gpu'`で定義できます（GPUの場合は追加で設定が必要になります[SageMakerのドキュメント](https://sagemaker.readthedocs.io/en/stable/overview.html#local-mode)を参照してください）。    
# **_Note：これはSageMaker Studio内では機能しません_**

# requirements.txtはトレーニングジョブの実行前に実行されます（コンテナにライブラリを追加する際に使用します）
# ファイルはここを参照しています。https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/requirements.txt
# 1点異なる部分は transformers >= 4.8.0 でHuggingFaceコンテナのバージョンが古く本家に追いついていないため、バージョンアップを行なっています。
# !pygmentize ./scripts/requirements.txt

# トレーニングジョブで実行されるコード
# 51行目の`check_min_version("4.10.0.dev0")`をコメントアウトした以外はSageMakerで実行するために変更を加えた部分はありません
# !pygmentize ./scripts/run_summarization.py

# +
from sagemaker.huggingface import HuggingFace

# hyperparameters, which are passed into the training job
hyperparameters={
    'model_name_or_path':'sonoisa/t5-base-japanese',
    'train_file': '/opt/ml/input/data/train/train.csv',
    'validation_file': '/opt/ml/input/data/validation/dev.csv',
    'test_file': '/opt/ml/input/data/test/test.csv',
    'text_column': 'src',
    'summary_column': 'tgt',
    'max_source_length': 512,
    'max_target_length': 64,
    'do_train': 'True',
    'do_eval': 'True',
    'do_predict': 'True',
    'predict_with_generate': 'True',
    'num_train_epochs': 5,
    'per_device_train_batch_size': 2,
    'per_device_eval_batch_size': 2,
    'use_fast_tokenizer': 'False',
    'save_steps': 500,
    'save_total_limit': 1,
    'output_dir':'/opt/ml/model',
}
# -

# ## Creating an Estimator and start a training job

# estimator
huggingface_estimator = HuggingFace(
    role=role,
    entry_point='run_summarization.py',
    source_dir='./scripts',
    instance_type='ml.p3.8xlarge',
    instance_count=1,
    volume_size=200,
    transformers_version='4.6',
    pytorch_version='1.7',
    py_version='py36',
    hyperparameters=hyperparameters,
)

# +
# starting the train job with our uploaded datasets as input
huggingface_estimator.fit({'train': input_train, 'validation': input_validation, 'test': input_test})

# ml.p3.8xlarge, 5 epochでの実行時間の目安
# Training seconds: 2894
# Billable seconds: 2894
# -

# ## Download-fine-tuned-model-from-s3
#

# +
import os

OUTPUT_DIR = './output/'
if not os.path.exists(OUTPUT_DIR):
    os.makedirs(OUTPUT_DIR)

# +
from sagemaker.s3 import S3Downloader

# 学習したモデルのダウンロード
S3Downloader.download(
    s3_uri=huggingface_estimator.model_data, # s3 uri where the trained model is located
    local_path='.', # local path where *.targ.gz is saved
    sagemaker_session=sess # sagemaker session used for training the model
)

# +
# OUTPUT_DIRに解凍します

# !tar -zxvf model.tar.gz -C output
# -

# ## Summarization on Local

# +
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained('output/')    
tokenizer = AutoTokenizer.from_pretrained('sonoisa/t5-base-japanese') 
model.eval()

# +
text = dev.src[1]
inputs = tokenizer.encode(text, return_tensors="pt", max_length=1024, truncation=True)

print('='*5,'原文', '='*5)
print(text)
print('-'*5, '要約', '-'*5)

with torch.no_grad():
    summary_ids = model.generate(inputs, max_length=1024, do_sample=False, num_beams=1)
    summary = tokenizer.decode(summary_ids[0])
    print(summary)
# -



StarcoderdataJupyterScriptsDedupFiltered