Init repo for TIABotV2

.gitattributes

@@ -32,3 +32,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model-sagemaker-5epochs/pytorch_model.bin filter=lfs diff=lfs merge=lfs -text
+train.csv filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,10 +1,10 @@
 ---
-title: TIABotV2
-emoji: 🐢
-colorFrom: purple
-colorTo: purple
+title: TIABot
+emoji: 🚀
+colorFrom: green
+colorTo: indigo
 sdk: gradio
-sdk_version: 3.9
+sdk_version: 3.1.1
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -0,0 +1,253 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[10]:
+
+
+import pandas as pd
+import os
+import torch
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+from transformers.optimization import  Adafactor 
+import time
+import warnings
+import random
+warnings.filterwarnings('ignore')
+
+import re
+
+def strip_html(text):
+    return re.sub('<[^<]+?>', '', text)
+
+
+# In[5]:
+
+
+train_columns = ['round_amount', 'round_date', 'stage', 'investee',
+       'investee_description', 'investee_country', 'investee_region',
+       'investee_subregion', 'investee_vertical', 'investee_industry',
+       'investor_list', 'previous_investors', 'prior_funding']
+train = pd.read_csv("train.csv")
+
+
+# In[6]:
+
+
+train.publication_timestamp = pd.to_datetime(train.publication_timestamp)
+
+
+# In[7]:
+
+
+input_text = train[train_columns].to_dict(orient='records')
+train_df = train[['title']].rename(columns={'title':'target_text'})
+train_df['input_text'] = input_text
+train_df['prefix'] = 'tia'
+train_df.input_text = train_df.input_text.astype(str)
+
+
+# In[8]:
+
+
+if torch.cuda.is_available():
+    dev = torch.device("cuda:0") 
+    print("Running on the GPU")
+else:
+    dev = torch.device("cpu")
+    print("Running on the CPU")
+
+
+# In[ ]:
+
+
+tokenizer = T5Tokenizer.from_pretrained('google/t5-v1_1-base')
+model = T5ForConditionalGeneration.from_pretrained('model-sagemaker-5epochs/', local_files_only=True)
+#moving the model to device(GPU/CPU)
+model.to(dev)
+
+
+# In[12]:
+
+
+vi_table = train[['investee_industry', 'investee_vertical']].drop_duplicates()
+
+
+# In[13]:
+
+
+def update_industry(value):
+    verticals = list(vi_table[vi_table['investee_industry'] == value]['investee_vertical'].values)
+    return verticals[0]
+
+def update_vertical(value):
+    industries = list(vi_table[vi_table['investee_vertical'] == value]['investee_industry'].values)
+    return industries[0]
+
+
+# In[ ]:
+
+
+update_industry('Green')
+
+
+# In[ ]:
+
+
+update_vertical('Clean tech')
+
+
+# In[ ]:
+
+
+import gradio as gr
+
+
+# In[ ]:
+
+
+num_return_sequences = 5
+
+
+# In[ ]:
+
+
+def generate_headline(stage, investee_country, investee_subregion, investee_region, 
+                      investee_vertical, investee_industry, 
+                      round_amount, investee, investee_description, investor_list, previous_investors,
+                      other_values):
+
+    full_df = other_values.set_index("key").T
+    
+    full_df['stage'] = stage
+    full_df['investee_country'] = investee_country
+    full_df['investee_subregion'] = investee_subregion
+    full_df['investee_region'] = investee_region
+    full_df['investee_vertical'] = investee_vertical
+    full_df['investee_industry'] = investee_industry
+    full_df['round_amount'] = str(float(round_amount))
+    full_df['investee'] = investee
+    full_df['investee_description'] = investee_description
+    full_df['investor_list'] = investor_list
+    full_df['previous_investors'] = previous_investors
+    
+    random_set =full_df[['round_amount', 'round_date', 'stage', 'investee',
+       'investee_description', 'investee_country', 'investee_region',
+       'investee_subregion', 'investee_vertical', 'investee_industry',
+       'investor_list', 'previous_investors', 'prior_funding']].to_json(orient="records")
+#    print(random_set)
+    
+    input_ids = tokenizer.encode(f"tia: {{{random_set}}}", return_tensors="pt")  # Batch size 1
+    input_ids=input_ids.to(dev)
+    outputs = model.generate(input_ids)
+    # text_output = tokenizer.decode(outputs[0]) # Single output
+    text_outputs = model.generate(inputs=input_ids, do_sample=True, 
+                                  num_beams=2, 
+                                  num_return_sequences=num_return_sequences,
+                                 repetition_penalty=5.0)
+    outputs = [strip_html(tokenizer.decode(o)) for o in text_outputs]
+    return "\n".join(outputs)
+
+
+# In[ ]:
+
+
+other_columns = ['round_date', 'prior_funding']
+
+
+# In[ ]:
+
+
+train.sample(1)[other_columns].T.reset_index().values
+
+
+# In[ ]:
+
+
+print(train.query("investee == 'NOSH'")['title'].head(1).T)
+train.query("investee == 'NOSH'")[train_columns].head(1).T
+
+
+# In[ ]:
+
+
+fake_data = {
+   "round_amount":1000000.0,
+   "round_date":"2018-09-26",
+   "stage":"Pre-series A",
+   "investee":"NOSH",
+   "investee_description":"NOSH makes and delivers ready-to-eat meals in Hong Kong.",
+   "investee_country":"Hong Kong",
+   "investee_region":"Asia",
+   "investee_subregion":"Eastern Asia",
+   "investee_vertical":"Food tech",
+   "investee_industry":"Restaurants & Food",
+   "investor_list":["Alibaba Entrepreneurs Fund (阿里巴巴创业者基金)"],
+   "previous_investors":"",
+   "prior_funding":1000000.0
+}
+
+
+# In[ ]:
+
+
+pd.DataFrame([fake_data]).T
+
+
+# In[ ]:
+
+
+demo = gr.Blocks()
+
+random_sample = train[train_columns].sample(1)
+random_sample = pd.DataFrame([fake_data])
+
+stage = gr.Dropdown(label="stage", choices=list(train[train_columns].stage.unique()))
+investee_country = gr.Dropdown(label="investee_country", choices=list(train[train_columns].investee_country.unique()),
+                              value=random_sample.investee_country.values[0])
+investee_subregion = gr.Dropdown(label="investee_subregion", choices=list(train[train_columns].investee_subregion.unique()),
+                              value=random_sample.investee_subregion.values[0])
+investee_region = gr.Dropdown(label="investee_region", choices=list(train[train_columns].investee_region.unique()),
+                             value=random_sample.investee_region.values[0])
+investee_vertical = gr.Dropdown(label="investee_vertical", choices=list(train[train_columns].investee_vertical.unique()),
+                            value=random_sample.investee_vertical.values[0])
+investee_industry = gr.Dropdown(label="investee_industry", choices=list(train[train_columns].investee_industry.unique()),
+                            value=random_sample.investee_industry.values[0])
+
+if pd.isnull(random_sample.round_amount.values[0]):
+    rand_amount = 0
+else:
+    rand_amount = random_sample.round_amount.values[0]
+
+round_amount = gr.Slider(label="round_amount", minimum=100000, maximum=200000000, 
+                         value=rand_amount, 
+                         step=100000)
+
+investee = gr.Textbox(label="investee", value=random_sample.investee.values[0])
+investee_description = gr.Textbox(label="investee_description", 
+           value=random_sample.investee_description.values[0])
+investor_list = gr.Textbox(label="investor_list", 
+           value=random_sample.investor_list.values[0])
+previous_investors = gr.Textbox(label="previous_investors", 
+           value=random_sample.previous_investors.values[0])  
+other_values = gr.Dataframe(
+    headers=['key', 'value'],
+    value=[['round_date', random_sample.round_date.values[0]],
+   ['prior_funding', random_sample.prior_funding.values[0]]]
+)
+out = gr.Textbox(max_lines=num_return_sequences)
+
+with demo:
+    gr.Markdown("Enter funding data to generate news headline.")
+   
+    inputs=[stage, investee_country, investee_subregion, investee_region, 
+                      investee_vertical, investee_industry, 
+                      round_amount, investee, investee_description, investor_list, previous_investors,
+                      other_values]
+    
+    investee_industry.change(fn=update_industry, inputs=investee_industry, outputs=investee_vertical)
+    investee_vertical.change(fn=update_vertical, inputs=investee_vertical, outputs=investee_industry)
+    gr.Interface(fn=generate_headline, inputs=inputs, outputs=out, live=True)
+    description="Enter funding data to generate news headline.",
+    live=True
+    
+demo.launch()

model-sagemaker-5epochs/code/.ipynb_checkpoints/inference-checkpoint.py

@@ -0,0 +1,50 @@
+# This is the script that will be used in the inference container
+import json
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def model_fn(model_dir):
+    """
+    Load the model and tokenizer for inference
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_dir).to(device).eval()
+
+    return {"model": model, "tokenizer": tokenizer}
+
+
+def predict_fn(input_data, model_dict):
+    """
+    Make a prediction with the model
+    """
+    text = input_data.pop("inputs")
+    parameters_list = input_data.pop("parameters_list", None)
+
+    tokenizer = model_dict["tokenizer"]
+    model = model_dict["model"]
+
+    # Parameters may or may not be passed
+    input_ids = tokenizer(
+        text, truncation=True, padding="longest", return_tensors="pt"
+    ).input_ids.to(device)
+
+    if parameters_list:
+        predictions = []
+        for parameters in parameters_list:
+            output = model.generate(input_ids, **parameters)
+            predictions.append(tokenizer.batch_decode(output, skip_special_tokens=True))
+    else:
+        output = model.generate(input_ids)
+        predictions = tokenizer.batch_decode(output, skip_special_tokens=True)
+
+    return predictions
+
+
+def input_fn(request_body, request_content_type):
+    """
+    Transform the input request to a dictionary
+    """
+    return json.loads(request_body)

model-sagemaker-5epochs/code/.ipynb_checkpoints/requirements-checkpoint.txt

@@ -0,0 +1,2 @@
+nltk
+rouge_score

model-sagemaker-5epochs/code/.ipynb_checkpoints/train-checkpoint.py

@@ -0,0 +1,215 @@
+# This is the script that will be used in the training container
+import argparse
+import logging
+import os
+import sys
+
+import numpy as np
+import nltk
+
+try:
+    nltk.data.find("tokenizers/punkt")
+except LookupError as e:
+    print(e)
+    try:
+        nltk.download("punkt")
+    except FileExistsError as e:
+        print(e)
+        pass
+
+from nltk import sent_tokenize
+
+from datasets import load_metric, load_from_disk
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+def tokenize(batch, text_column, target_column, max_source, max_target):
+    tokenized_input = tokenizer(
+        batch[text_column], padding="max_length", truncation=True, max_length=max_source
+    )
+    tokenized_target = tokenizer(
+        batch[target_column],
+        padding="max_length",
+        truncation=True,
+        max_length=max_target,
+    )
+
+    tokenized_input["labels"] = tokenized_target["input_ids"]
+
+    return tokenized_input
+
+
+def load_and_tokenize_dataset(
+    data_dir, split, text_column, target_column, max_source, max_target
+):
+
+    dataset = load_from_disk(os.path.join(data_dir, split))
+    tokenized_dataset = dataset.map(
+        lambda x: tokenize(x, text_column, target_column, max_source, max_target),
+        batched=True,
+        batch_size=512,
+    )
+    tokenized_dataset.set_format(
+        "numpy", columns=["input_ids", "attention_mask", "labels"]
+    )
+
+    return tokenized_dataset
+
+def compute_metrics(eval_pred):
+    metric = load_metric('glue', 'mrpc')
+    predictions, references = eval_pred
+    return metric.compute(predictions=predictions, references=references)
+    
+#     metric = load_metric("rouge")
+#     predictions, labels = eval_pred
+#     decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+    
+#     # Replace -100 in the labels as we can't decode them.
+#     labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+#     decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+    
+#     # Rouge expects a newline after each sentence
+#     decoded_preds = ["\n".join(nltk.sent_tokenize(pred.strip()))
+#                       for pred in decoded_preds]
+#     decoded_labels = ["\n".join(nltk.sent_tokenize(label.strip())) 
+#                       for label in decoded_labels]
+    
+#     # Compute ROUGE scores
+#     logger.info("Decoded preds: %s" % decoded_preds)
+#     logger.info("Decoded labels: %s" % decoded_labels)
+#     result = metric.compute(predictions=decoded_preds, references=decoded_labels,
+#                             use_stemmer=True)
+
+#     # Extract ROUGE f1 scores
+#     result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
+    
+#     # Add mean generated length to metrics
+#     prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id)
+#                       for pred in predictions]
+#     result["gen_len"] = np.mean(prediction_lens)
+    
+#     return {k: round(v, 4) for k, v in result.items()}
+
+def train(args):
+    from transformers import T5ForConditionalGeneration, T5Tokenizer
+    
+    logger.info("Loading tokenizer...\n")
+    global tokenizer
+    global model_name
+    model_name = args.model_name
+    
+    logger.info("Loading pretrained model\n")
+    if "google" in model_name:
+        model = T5ForConditionalGeneration.from_pretrained(model_name)    
+        tokenizer = T5Tokenizer.from_pretrained(model_name)
+    else:
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+    logger.info("Pretrained model loaded\n")
+
+    logger.info("Fetching and tokenizing data for training")
+    train_dataset = load_and_tokenize_dataset(
+        args.train_data_dir,
+        "train",
+        args.text_column,
+        args.target_column,
+        args.max_source,
+        args.max_target,
+    )
+
+    logger.info("Tokenizing data for training loaded")
+
+    eval_dataset = load_and_tokenize_dataset(
+        args.train_data_dir,
+        "validation",
+        args.text_column,
+        args.target_column,
+        args.max_source,
+        args.max_target,
+    )
+    test_dataset = load_and_tokenize_dataset(
+        args.train_data_dir,
+        "test",
+        args.text_column,
+        args.target_column,
+        args.max_source,
+        args.max_target,
+    )
+
+    logger.info("Defining training arguments\n")
+    training_args = Seq2SeqTrainingArguments(
+        output_dir=args.model_dir,
+        num_train_epochs=args.epoch,
+        per_device_train_batch_size=args.train_batch_size,
+        per_device_eval_batch_size=args.eval_batch_size,
+        learning_rate=args.lr,
+        warmup_steps=args.warmup_steps,
+        weight_decay=args.weight_decay,
+        logging_dir=args.log_dir,
+        logging_strategy=args.logging_strategy,
+        load_best_model_at_end=True,
+        adafactor=True,
+        do_train=True,
+        do_eval=True,
+        do_predict=True,
+        save_total_limit=3,
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        predict_with_generate=True,
+        metric_for_best_model="eval_loss",
+        seed=7,
+    )
+
+    logger.info("Defining seq2seq Trainer")
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        tokenizer=tokenizer#,compute_metrics=compute_metrics,
+    )
+
+    logger.info("Starting Training")
+    trainer.train()
+    logger.info("Model trained successfully")
+    trainer.save_model()
+    logger.info("Model saved successfully")
+
+    # Evaluation
+    logger.info("*** Evaluate on test set***")
+
+    logger.info(trainer.predict(test_dataset))
+
+    logger.info("Removing unused checkpoints to save space in container")
+    os.system(f"rm -rf {args.model_dir}/checkpoint-*/")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-name", type=str, default="google/pegasus-xsum")
+    parser.add_argument(
+        "--train-data-dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"]
+    )
+    # parser.add_argument("--val-data-dir", type=str,
+    #                   default=os.environ["SM_CHANNEL_VALIDATION"])
+    # parser.add_argument("--test-data-dir", type=str,
+    #                    default=os.environ["SM_CHANNEL_TEST"])
+    parser.add_argument("--text-column", type=str, default="dialogue")
+    parser.add_argument("--target-column", type=str, default="summary")
+    parser.add_argument("--max-source", type=int, default=512)
+    parser.add_argument("--max-target", type=int, default=80)
+    parser.add_argument("--model-dir", type=str, default=os.environ["SM_MODEL_DIR"])
+    parser.add_argument("--epoch", type=int, default=5)
+    parser.add_argument("--train-batch-size", type=int, default=2)
+    parser.add_argument("--eval-batch-size", type=int, default=2)
+    parser.add_argument("--warmup-steps", type=float, default=500)
+    parser.add_argument("--lr", type=float, default=2e-5)
+    parser.add_argument("--weight-decay", type=float, default=0.0)
+    parser.add_argument("--log-dir", type=str, default=os.environ["SM_OUTPUT_DIR"])
+    parser.add_argument("--logging-strategy", type=str, default="epoch")
+    train(parser.parse_args())

model-sagemaker-5epochs/code/inference.py

@@ -0,0 +1,50 @@
+# This is the script that will be used in the inference container
+import json
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def model_fn(model_dir):
+    """
+    Load the model and tokenizer for inference
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_dir).to(device).eval()
+
+    return {"model": model, "tokenizer": tokenizer}
+
+
+def predict_fn(input_data, model_dict):
+    """
+    Make a prediction with the model
+    """
+    text = input_data.pop("inputs")
+    parameters_list = input_data.pop("parameters_list", None)
+
+    tokenizer = model_dict["tokenizer"]
+    model = model_dict["model"]
+
+    # Parameters may or may not be passed
+    input_ids = tokenizer(
+        text, truncation=True, padding="longest", return_tensors="pt"
+    ).input_ids.to(device)
+
+    if parameters_list:
+        predictions = []
+        for parameters in parameters_list:
+            output = model.generate(input_ids, **parameters)
+            predictions.append(tokenizer.batch_decode(output, skip_special_tokens=True))
+    else:
+        output = model.generate(input_ids)
+        predictions = tokenizer.batch_decode(output, skip_special_tokens=True)
+
+    return predictions
+
+
+def input_fn(request_body, request_content_type):
+    """
+    Transform the input request to a dictionary
+    """
+    return json.loads(request_body)

model-sagemaker-5epochs/code/requirements.txt

@@ -0,0 +1,2 @@
+nltk
+rouge_score

model-sagemaker-5epochs/code/train.py

@@ -0,0 +1,215 @@
+# This is the script that will be used in the training container
+import argparse
+import logging
+import os
+import sys
+
+import numpy as np
+import nltk
+
+try:
+    nltk.data.find("tokenizers/punkt")
+except LookupError as e:
+    print(e)
+    try:
+        nltk.download("punkt")
+    except FileExistsError as e:
+        print(e)
+        pass
+
+from nltk import sent_tokenize
+
+from datasets import load_metric, load_from_disk
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+def tokenize(batch, text_column, target_column, max_source, max_target):
+    tokenized_input = tokenizer(
+        batch[text_column], padding="max_length", truncation=True, max_length=max_source
+    )
+    tokenized_target = tokenizer(
+        batch[target_column],
+        padding="max_length",
+        truncation=True,
+        max_length=max_target,
+    )
+
+    tokenized_input["labels"] = tokenized_target["input_ids"]
+
+    return tokenized_input
+
+
+def load_and_tokenize_dataset(
+    data_dir, split, text_column, target_column, max_source, max_target
+):
+
+    dataset = load_from_disk(os.path.join(data_dir, split))
+    tokenized_dataset = dataset.map(
+        lambda x: tokenize(x, text_column, target_column, max_source, max_target),
+        batched=True,
+        batch_size=512,
+    )
+    tokenized_dataset.set_format(
+        "numpy", columns=["input_ids", "attention_mask", "labels"]
+    )
+
+    return tokenized_dataset
+
+def compute_metrics(eval_pred):
+    metric = load_metric('glue', 'mrpc')
+    predictions, references = eval_pred
+    return metric.compute(predictions=predictions, references=references)
+    
+#     metric = load_metric("rouge")
+#     predictions, labels = eval_pred
+#     decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+    
+#     # Replace -100 in the labels as we can't decode them.
+#     labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+#     decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+    
+#     # Rouge expects a newline after each sentence
+#     decoded_preds = ["\n".join(nltk.sent_tokenize(pred.strip()))
+#                       for pred in decoded_preds]
+#     decoded_labels = ["\n".join(nltk.sent_tokenize(label.strip())) 
+#                       for label in decoded_labels]
+    
+#     # Compute ROUGE scores
+#     logger.info("Decoded preds: %s" % decoded_preds)
+#     logger.info("Decoded labels: %s" % decoded_labels)
+#     result = metric.compute(predictions=decoded_preds, references=decoded_labels,
+#                             use_stemmer=True)
+
+#     # Extract ROUGE f1 scores
+#     result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
+    
+#     # Add mean generated length to metrics
+#     prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id)
+#                       for pred in predictions]
+#     result["gen_len"] = np.mean(prediction_lens)
+    
+#     return {k: round(v, 4) for k, v in result.items()}
+
+def train(args):
+    from transformers import T5ForConditionalGeneration, T5Tokenizer
+    
+    logger.info("Loading tokenizer...\n")
+    global tokenizer
+    global model_name
+    model_name = args.model_name
+    
+    logger.info("Loading pretrained model\n")
+    if "google" in model_name:
+        model = T5ForConditionalGeneration.from_pretrained(model_name)    
+        tokenizer = T5Tokenizer.from_pretrained(model_name)
+    else:
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+    logger.info("Pretrained model loaded\n")
+
+    logger.info("Fetching and tokenizing data for training")
+    train_dataset = load_and_tokenize_dataset(
+        args.train_data_dir,
+        "train",
+        args.text_column,
+        args.target_column,
+        args.max_source,
+        args.max_target,
+    )
+
+    logger.info("Tokenizing data for training loaded")
+
+    eval_dataset = load_and_tokenize_dataset(
+        args.train_data_dir,
+        "validation",
+        args.text_column,
+        args.target_column,
+        args.max_source,
+        args.max_target,
+    )
+    test_dataset = load_and_tokenize_dataset(
+        args.train_data_dir,
+        "test",
+        args.text_column,
+        args.target_column,
+        args.max_source,
+        args.max_target,
+    )
+
+    logger.info("Defining training arguments\n")
+    training_args = Seq2SeqTrainingArguments(
+        output_dir=args.model_dir,
+        num_train_epochs=args.epoch,
+        per_device_train_batch_size=args.train_batch_size,
+        per_device_eval_batch_size=args.eval_batch_size,
+        learning_rate=args.lr,
+        warmup_steps=args.warmup_steps,
+        weight_decay=args.weight_decay,
+        logging_dir=args.log_dir,
+        logging_strategy=args.logging_strategy,
+        load_best_model_at_end=True,
+        adafactor=True,
+        do_train=True,
+        do_eval=True,
+        do_predict=True,
+        save_total_limit=3,
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        predict_with_generate=True,
+        metric_for_best_model="eval_loss",
+        seed=7,
+    )
+
+    logger.info("Defining seq2seq Trainer")
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        tokenizer=tokenizer#,compute_metrics=compute_metrics,
+    )
+
+    logger.info("Starting Training")
+    trainer.train()
+    logger.info("Model trained successfully")
+    trainer.save_model()
+    logger.info("Model saved successfully")
+
+    # Evaluation
+    logger.info("*** Evaluate on test set***")
+
+    logger.info(trainer.predict(test_dataset))
+
+    logger.info("Removing unused checkpoints to save space in container")
+    os.system(f"rm -rf {args.model_dir}/checkpoint-*/")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-name", type=str, default="google/pegasus-xsum")
+    parser.add_argument(
+        "--train-data-dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"]
+    )
+    # parser.add_argument("--val-data-dir", type=str,
+    #                   default=os.environ["SM_CHANNEL_VALIDATION"])
+    # parser.add_argument("--test-data-dir", type=str,
+    #                    default=os.environ["SM_CHANNEL_TEST"])
+    parser.add_argument("--text-column", type=str, default="dialogue")
+    parser.add_argument("--target-column", type=str, default="summary")
+    parser.add_argument("--max-source", type=int, default=512)
+    parser.add_argument("--max-target", type=int, default=80)
+    parser.add_argument("--model-dir", type=str, default=os.environ["SM_MODEL_DIR"])
+    parser.add_argument("--epoch", type=int, default=5)
+    parser.add_argument("--train-batch-size", type=int, default=2)
+    parser.add_argument("--eval-batch-size", type=int, default=2)
+    parser.add_argument("--warmup-steps", type=float, default=500)
+    parser.add_argument("--lr", type=float, default=2e-5)
+    parser.add_argument("--weight-decay", type=float, default=0.0)
+    parser.add_argument("--log-dir", type=str, default=os.environ["SM_OUTPUT_DIR"])
+    parser.add_argument("--logging-strategy", type=str, default="epoch")
+    train(parser.parse_args())

model-sagemaker-5epochs/config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "google/t5-v1_1-base",
+  "architectures": [
+    "T5ForConditionalGeneration"
+  ],
+  "d_ff": 2048,
+  "d_kv": 64,
+  "d_model": 768,
+  "decoder_start_token_id": 0,
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "gradient_checkpointing": false,
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "t5",
+  "num_decoder_layers": 12,
+  "num_heads": 12,
+  "num_layers": 12,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_num_buckets": 32,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.6.1",
+  "use_cache": true,
+  "vocab_size": 32128
+}

model-sagemaker-5epochs/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d13160a831c6a1f94d5205c958f8b79a9468ed83497ed3098aeb7645d0d109a1
+size 990445401

model-sagemaker-5epochs/special_tokens_map.json

@@ -0,0 +1 @@
+{"eos_token": "</s>", "unk_token": "<unk>", "pad_token": "<pad>", "additional_special_tokens": ["<extra_id_0>", "<extra_id_1>", "<extra_id_2>", "<extra_id_3>", "<extra_id_4>", "<extra_id_5>", "<extra_id_6>", "<extra_id_7>", "<extra_id_8>", "<extra_id_9>", "<extra_id_10>", "<extra_id_11>", "<extra_id_12>", "<extra_id_13>", "<extra_id_14>", "<extra_id_15>", "<extra_id_16>", "<extra_id_17>", "<extra_id_18>", "<extra_id_19>", "<extra_id_20>", "<extra_id_21>", "<extra_id_22>", "<extra_id_23>", "<extra_id_24>", "<extra_id_25>", "<extra_id_26>", "<extra_id_27>", "<extra_id_28>", "<extra_id_29>", "<extra_id_30>", "<extra_id_31>", "<extra_id_32>", "<extra_id_33>", "<extra_id_34>", "<extra_id_35>", "<extra_id_36>", "<extra_id_37>", "<extra_id_38>", "<extra_id_39>", "<extra_id_40>", "<extra_id_41>", "<extra_id_42>", "<extra_id_43>", "<extra_id_44>", "<extra_id_45>", "<extra_id_46>", "<extra_id_47>", "<extra_id_48>", "<extra_id_49>", "<extra_id_50>", "<extra_id_51>", "<extra_id_52>", "<extra_id_53>", "<extra_id_54>", "<extra_id_55>", "<extra_id_56>", "<extra_id_57>", "<extra_id_58>", "<extra_id_59>", "<extra_id_60>", "<extra_id_61>", "<extra_id_62>", "<extra_id_63>", "<extra_id_64>", "<extra_id_65>", "<extra_id_66>", "<extra_id_67>", "<extra_id_68>", "<extra_id_69>", "<extra_id_70>", "<extra_id_71>", "<extra_id_72>", "<extra_id_73>", "<extra_id_74>", "<extra_id_75>", "<extra_id_76>", "<extra_id_77>", "<extra_id_78>", "<extra_id_79>", "<extra_id_80>", "<extra_id_81>", "<extra_id_82>", "<extra_id_83>", "<extra_id_84>", "<extra_id_85>", "<extra_id_86>", "<extra_id_87>", "<extra_id_88>", "<extra_id_89>", "<extra_id_90>", "<extra_id_91>", "<extra_id_92>", "<extra_id_93>", "<extra_id_94>", "<extra_id_95>", "<extra_id_96>", "<extra_id_97>", "<extra_id_98>", "<extra_id_99>"]}

model-sagemaker-5epochs/spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d60acb128cf7b7f2536e8f38a5b18a05535c9e14c7a355904270e15b0945ea86
+size 791656

model-sagemaker-5epochs/tokenizer_config.json

@@ -0,0 +1 @@
+{"eos_token": "</s>", "unk_token": "<unk>", "pad_token": "<pad>", "extra_ids": 100, "additional_special_tokens": ["<extra_id_0>", "<extra_id_1>", "<extra_id_2>", "<extra_id_3>", "<extra_id_4>", "<extra_id_5>", "<extra_id_6>", "<extra_id_7>", "<extra_id_8>", "<extra_id_9>", "<extra_id_10>", "<extra_id_11>", "<extra_id_12>", "<extra_id_13>", "<extra_id_14>", "<extra_id_15>", "<extra_id_16>", "<extra_id_17>", "<extra_id_18>", "<extra_id_19>", "<extra_id_20>", "<extra_id_21>", "<extra_id_22>", "<extra_id_23>", "<extra_id_24>", "<extra_id_25>", "<extra_id_26>", "<extra_id_27>", "<extra_id_28>", "<extra_id_29>", "<extra_id_30>", "<extra_id_31>", "<extra_id_32>", "<extra_id_33>", "<extra_id_34>", "<extra_id_35>", "<extra_id_36>", "<extra_id_37>", "<extra_id_38>", "<extra_id_39>", "<extra_id_40>", "<extra_id_41>", "<extra_id_42>", "<extra_id_43>", "<extra_id_44>", "<extra_id_45>", "<extra_id_46>", "<extra_id_47>", "<extra_id_48>", "<extra_id_49>", "<extra_id_50>", "<extra_id_51>", "<extra_id_52>", "<extra_id_53>", "<extra_id_54>", "<extra_id_55>", "<extra_id_56>", "<extra_id_57>", "<extra_id_58>", "<extra_id_59>", "<extra_id_60>", "<extra_id_61>", "<extra_id_62>", "<extra_id_63>", "<extra_id_64>", "<extra_id_65>", "<extra_id_66>", "<extra_id_67>", "<extra_id_68>", "<extra_id_69>", "<extra_id_70>", "<extra_id_71>", "<extra_id_72>", "<extra_id_73>", "<extra_id_74>", "<extra_id_75>", "<extra_id_76>", "<extra_id_77>", "<extra_id_78>", "<extra_id_79>", "<extra_id_80>", "<extra_id_81>", "<extra_id_82>", "<extra_id_83>", "<extra_id_84>", "<extra_id_85>", "<extra_id_86>", "<extra_id_87>", "<extra_id_88>", "<extra_id_89>", "<extra_id_90>", "<extra_id_91>", "<extra_id_92>", "<extra_id_93>", "<extra_id_94>", "<extra_id_95>", "<extra_id_96>", "<extra_id_97>", "<extra_id_98>", "<extra_id_99>"], "model_max_length": 512, "name_or_path": "google/t5-v1_1-base", "special_tokens_map_file": "/root/.cache/huggingface/transformers/76bf19bfedb85afbe644966ca9ab7b0404d753a41bf601115bced39f825ffa9c.c94798918c92ded6aeef2d2f0e666d2cc4145eca1aa6e1336fde07f2e13e2f46", "tokenizer_file": null}

model-sagemaker-5epochs/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92e99d1bdd152d96b8c6dd0796be148e20b4a351b909b0744df5c127c8709351
+size 2479

requirements.txt

@@ -0,0 +1,6 @@
+gradio==3.0.22
+pandas==1.4.2
+torch==1.10.2
+transformers==4.19.2
+sentencepiece==0.1.95
+sentence-transformers==2.2.2

train.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:999febaa4d0e013cb0c89ba43c657bfdf13d9d7d8e52f4050b64341ff833489d
+size 34332589