Upload 29 files

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Dinh Truong Phan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,12 +1,109 @@
----
-title: Dialogue Text Summarization
-emoji: ⚡
-colorFrom: yellow
-colorTo: blue
-sdk: streamlit
-sdk_version: 1.35.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Problem Description
+
+This project aims to develop a system capable of automatically **summarizing short dialogue text**. This addresses the challenge of extracting concise yet informative summaries from conversational exchanges, enabling users to **quickly grasp the information of the dialogues**.
+
+Summarizing these conversations can be valuable for various applications, such as:
+- Streamlining information retrieval in customer service interactions
+- Condensing meeting discussions for efficient review
+- Providing concise overviews of chat conversations on social media platforms
+
+This project tackles the task of automatically generating concise summaries, saving users time and effort while improving comprehension.
+
+![](assets/image2.png)
+
+<p align="center"><i>Source: Google Research</i></p>
+
+**Input:** Dialogue text
+
+Example: 
+```
+Matt: Do you want to go for date? 
+Agnes: Wow! You caught me out with this question Matt. 
+...
+Agnes: See you on saturday. 
+Matt: Yes, looking forward to it. 
+Agnes: Me too.
+```
+
+**Output:** Summarized dialogue
+
+Example: 
+```
+Matt invites Agnes for a date to get to know each other better. They'll go to the Georgian restaurant in Kazimierz on Saturday at 6 pm, and he'll pick her up on the way to the place.
+```
+
+# Dataset
+
+We'll utilize the `DialogSum` dataset accessible from 🤗**Hugging Face** (https://huggingface.co/datasets/knkarthick/dialogsum) and **Paper** (https://arxiv.org/pdf/2105.06762.pdf). This dataset comprises real-life dialogue scenarios paired with corresponding manually crafted summaries and dialogue topics.
+
+`DialogSum` is a large-scale dialogue summarization dataset, consisting of **13,460** (Plus 100 holdout data for topic generation) dialogues with corresponding manually labeled summaries and topics.
+
+Here's a sample of the `DialogSum` dataset structure:
+
+
+|id|dialogue|summary|topic|
+|-|-|-|-|
+|train_3|#Person1#: Why didn't you tell me you had a girlfriend? #Person2#: Sorry, I thought you knew. ... #Person1#: Oh, you men! You are all the same.|#Person1#'s angry because #Person2# didn't tell #Person1# that #Person2# had a girlfriend and would marry her.|have a girl friend|
+|train_16|#Person1#: Tell me something about your Valentine's Day. ...#Person2#: Yeah, that is what the holiday is for, isn't it?|#Person2# tells #Person1# their Valentine's Day. #Person1# feels it's romantic.|Valentine's Day|
+|...|...|...|...|
+
+**Distribution of dataset**
+
+|Dialogue|Summary|Dialogue + Summary|
+|:-:|:-:|:-:|
+|![](assets/hist_dialogue.png)|![](assets/hist_summary.png)|![](assets/hist_dialogue+summary.png)|
+
+# Method
+
+### Pre-trained Language Models:
+
+This project explores two powerful LLMs well-suited for dialogue summarization:
+
+- **FLAN-T5:** This model excels at understanding complex relationships within text, making it effective in summarizing the nuances of conversations.
+- **BART:** This model boasts strong capabilities in text generation tasks, making it adept at generating informative and well-structured summaries.
+
+### Fine-tuning Techniques:
+
+To tailor these LLMs specifically for dialogue summarization, we will investigate several fine-tuning approaches:
+
+- Instruction Fine-tuning
+- Parameter Efficient Fine Tuning (PEFT)
+    + Low-Rank Adaptation **(LoRA)**
+    + Quantized Low-Rank Adaptation **(QLoRA)**
+
+# Installation
+
+```
+!git clone "https://github.com/dtruong46me/dialogue-text-summarization.git"
+```
+
+# Contributions
+
+**Supervisor:** Prof. Le Thanh Huong
+
+**Student Group:**
+
+|No.|Name|Student ID|Email|
+|:-:|-|:-:|-|
+|1|Phan Dinh Truong (Leader)|20214937|[email protected]|
+|2|Nguyen Tung Luong|20214913|[email protected]|
+|3|Vu Tuan Minh|20210597|[email protected]|
+|4|Hoang Tu Quyen|20214929|[email protected]|
+
+# [Bonus] How to run Streamlit on Kaggle
+
+```
+!pip install -q streamlit
+```
+
+```
+!wget -q -O - ipv4.icanhazip.com
+```
+
+```
+!npm install -g localtunnel -q
+```
+
+```
+!streamlit run "/kaggle/working/dialogue-text-summarization/streamlit_app.py" & npx localtunnel --port 8501
+```

app.py

@@ -0,0 +1,73 @@
+import streamlit as st
+
+from transformers import GenerationConfig, BartModel, BartTokenizer, AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+import time
+
+import sys, os
+
+path = os.path.abspath(os.path.dirname(__file__))
+sys.path.insert(0, path)
+
+from gen_summary import generate_summary
+
+
+st.title("Dialogue Text Summarization")
+st.caption("Natural Language Processing Project 20232")
+
+st.write("---")
+
+with st.sidebar:
+    checkpoint = st.selectbox("Model", options=[
+        "Choose model",
+        "dtruong46me/train-bart-base",
+        "dtruong46me/flant5-small",
+        "dtruong46me/flant5-base",
+        "dtruong46me/flan-t5-s",
+        "ntluongg/bart-base-luong"
+    ])
+    st.button("Model detail", use_container_width=True)
+    st.write("-----")
+    st.write("**Generate Options:**")
+    min_new_tokens = st.number_input("Min new tokens", min_value=1, max_value=64, value=10)
+    max_new_tokens = st.number_input("Max new tokens", min_value=64, max_value=128, value=64)
+    temperature = st.number_input("Temperature", min_value=0.0, max_value=1.0, value=0.9, step=0.05)
+    top_k = st.number_input("Top_k", min_value=1, max_value=50, step=1, value=20)
+    top_p = st.number_input("Top_p", min_value=0.01, max_value=1.00, step=0.01, value=1.0)
+
+
+height = 200
+
+input_text = st.text_area("Dialogue", height=height)
+
+generation_config = GenerationConfig(
+    min_new_tokens=min_new_tokens,
+    max_new_tokens=320,
+    temperature=temperature,
+    top_p=top_p,
+    top_k=top_k
+)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+if checkpoint=="Choose model":
+    tokenizer = None
+    model = None
+
+if checkpoint!="Choose model":
+    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+    model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
+
+
+
+if st.button("Submit"):
+    st.write("---")
+    st.write("## Summary")
+
+    if checkpoint=="Choose model":
+        st.error("Please selece a model!")
+
+    else:
+        if input_text=="":
+            st.error("Please enter a dialogue!")
+        st.write(generate_summary(model, " ".join(input_text.split()), generation_config, tokenizer))

gen_summary.py

@@ -0,0 +1,66 @@
+import torch
+from transformers import AutoTokenizer, GenerationConfig, TextStreamer, AutoModelForSeq2SeqLM
+
+import logging
+
+import warnings
+warnings.filterwarnings("ignore")
+
+# =  =  =  =  =  =  =  =  =  =  =  Logging Setup  =  =  =  =  =  =  =  =  =  =  =  =  = 
+logger = logging.getLogger(__name__)
+logging.basicConfig(
+    format  = "%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt = "%m/%d/%Y %H:%M:%S",
+    level   = logging.INFO,
+)
+# =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  = 
+
+def generate_summary(model, input_text, generation_config, tokenizer, st_container=None) -> str:
+
+    try:
+        prefix = "Summarize the following conversation: \n###\n"
+        suffix = "\n### Summary:"
+
+        input_ids = tokenizer.encode(prefix + input_text + "The generated summary should be around " + str(0.15*len(input_text)) + " words." + suffix, return_tensors="pt")
+        output_ids = model.generate(input_ids, do_sample=True, generation_config=generation_config)
+
+        if "bart" in model.name_or_path and model.name_or_path != "dtruong46me/bart-base-qds":
+            output_ids[0][1] = 2
+        
+        # streamer = TextStreamer(tokenizer, skip_special_tokens=True)
+        # model.generate(input_ids, streamer=streamer, do_sample=True, decoder_start_token_id=2, generation_config=generation_config)
+        # logger.info("\nComplete generate summary!")
+        output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        return output_text
+    
+    except Exception as e:
+        print(f"Error while generating: {e}")
+        raise e
+    
+if __name__=="__main__":
+    input = "#Person1#: Ms. Dawson, I need you to take a dictation for me. #Person2#: Yes, sir... #Person1#: This should go out as an intra-office memorandum to all employees by this afternoon. Are you ready? #Person2#: Yes, sir. Go ahead. #Person1#: Attention all staff... Effective immediately, all office communications are restricted to email correspondence and official memos. The use of Instant Message programs by employees during working hours is strictly prohibited. #Person2#: Sir, does this apply to intra-office communications only? Or will it also restrict external communications? #Person1#: It should apply to all communications, not only in this office between employees, but also any outside communications. #Person2#: But sir, many employees use Instant Messaging to communicate with their clients. #Person1#: They will just have to change their communication methods. I don't want any - one using Instant Messaging in this office. It wastes too much time! Now, please continue with the memo. Where were we? #Person2#: This applies to internal and external communications. #Person1#: Yes. Any employee who persists in using Instant Messaging will first receive a warning and be placed on probation. At second offense, the employee will face termination. Any questions regarding this new policy may be directed to department heads. #Person2#: Is that all? #Person1#: Yes. Please get this memo typed up and distributed to all employees before 4 pm."
+    target1 = "Ms. Dawson helps #Person1# to write a memo to inform every employee that they have to change the communication method and should not use Instant Messaging anymore."
+    target2 = "In order to prevent employees from wasting time on Instant Message programs, #Person1# decides to terminate the use of those programs and asks Ms. Dawson to send out a memo to all employees by the afternoon."
+    target3 = "Ms. Dawson takes a dictation for #Person1# about prohibiting the use of Instant Message programs in the office. They argue about its reasonability but #Person1# still insists."
+
+    generation_config = GenerationConfig(
+        min_new_tokens=10,
+        max_new_tokens=256,
+        temperature=0.9,
+        top_p=1.0,
+        top_k=50        
+    )
+
+    checkpoint = "dtruong46me/bart-base-qds2"
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+
+    model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint).to(device)
+    
+    generate_summary(model, input, generation_config, tokenizer)
+    print("\n==============\n")
+
+    print("Human base line:\n", target1, end="\n\n")
+    print("Human base line:\n", target2, end="\n\n")
+    print("Human base line:\n", target3, end="\n\n")

requirements.txt

@@ -0,0 +1,14 @@
+datasets
+huggingface_hub
+nltk
+numpy
+pandas
+peft
+replicate
+streamlit
+torch
+transformers==4.36.1
+wandb
+evaluate
+rouge_score
+bert_score

results/rouge_score.csv

@@ -0,0 +1,9 @@
+rouge1,rouge2,rougeL,rougeLsum,gen_len,checkpoint
+0.39233350039050524,0.1331263872944557,0.30561232240272806,0.305581876074012,25.568,dtruong46me/flant5-small
+0.42773411047439297,0.16070313389865537,0.33964372087731554,0.33971528751465496,24.633333333333333,dtruong46me/flant5-base
+0.4436612424628238,0.18215770435271772,0.3574836391515892,0.3575112795473217,25.358,dtruong46me/train-bart-base
+0.44596490799011734,0.1791041702437794,0.36099829444161424,0.3612203644902555,18.72,dtruong46me/bart-base-instructds2
+0.5335,0.2672,0.5084,0,0,human-annotated-summary
+0.4728,0.2118,0.4483,0,0,bart-large-in-paper
+0.5165,0.2981,0.4336,0.4337,23.187,dtruong46me/bart-base-qds
+0.4061788843274445,0.1588224274185049,0.3175643149646888,0.3207910509892517,26.058,dtruong46me/flan-t5-s

run_evaluation.py

@@ -0,0 +1,77 @@
+import warnings
+warnings.filterwarnings("ignore")
+
+from datasets import load_dataset
+
+import os, sys
+
+import pandas as pd
+import argparse
+
+path = os.path.abspath(os.path.join(os.path.dirname(__file__)))
+sys.path.insert(0, path)
+
+from src.model.model import load_model
+from src.evaluate.evaluation import evaluation_rouge
+from transformers import GenerationConfig
+
+
+def save_metrics_to_csv(results, resultpath, checkpoint):
+    
+    results["checkpoint"] = checkpoint
+
+    # Convert results to DataFrame
+    df = pd.DataFrame([results])
+
+    if not os.path.isfile(resultpath):
+        df.to_csv(resultpath, index=False)
+    else:
+        df.to_csv(resultpath, mode='a', header=False, index=False)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluation metric")
+    parser.add_argument("--datapath", type=str, default="knkarthick/dialogsum")
+    parser.add_argument("--checkpoint", type=str, default="google/flan-t5-base")
+    parser.add_argument("--resultpath", type=str, default="results/rouge_score.csv")
+    
+    parser.add_argument("--min_new_tokens", type=int, default=10)
+    parser.add_argument("--max_new_tokens", type=int, default=256)
+    parser.add_argument("--temperature", type=float, default=0.9)
+    parser.add_argument("--top_p", type=float, default=1.0)
+    parser.add_argument("--top_k", type=int, default=50)
+
+    args = parser.parse_args()
+
+    print("=========================================")
+    print('\n'.join(f' + {k}={v}' for k, v in vars(args).items()))
+    print("=========================================")
+    
+    datapath = args.datapath
+    checkpoint = args.checkpoint
+
+    generation_config = GenerationConfig(
+        min_new_tokens=args.min_new_tokens,
+        max_new_tokens=args.max_new_tokens,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        top_k=args.top_k
+    )
+
+    data = load_dataset("binwang/InstructDS_datasets", "DialogSum", split="test")
+
+    model = load_model(checkpoint)
+    print(f"Loaded model from: {checkpoint}")
+
+    results = evaluation_rouge(model, data, generation_config)
+    
+    print("--------------------------")
+    for k, v in results.items():
+        print(f"{k}: {v}")
+    print("--------------------------")
+
+    save_metrics_to_csv(results, args.resultpath, checkpoint)
+    print(f"Results saved to: {args.resultpath}")
+
+if __name__ == "__main__":
+    main()

run_training.py

@@ -0,0 +1,39 @@
+import wandb
+from huggingface_hub import login
+
+import warnings
+warnings.filterwarnings("ignore")
+
+import os
+import sys
+
+path = os.path.abspath(os.path.join(os.path.dirname(__file__)))
+sys.path.insert(0, path)
+
+from src.pipelines.training_pipeline import training_pipeline
+from src.utils import parse_args
+
+def main():
+    # Load argument parser
+    args = parse_args()
+    print(f"\033[92mLoaded argument parsers\033[00m")
+
+    # Load token ID
+    huggingface_hub_token = args.huggingface_hub_token
+    wandb_token = args.wandb_token
+
+    if wandb_token:
+        os.environ["WANDB_PROJECT"] = "nlp_project"
+
+    # Login to Huggingface Hub and WandB
+    login(token=huggingface_hub_token)
+    print("\033[92mSuccessful login to Huggingface Hub\033[00m")
+    
+    wandb.login(key=wandb_token)
+    print("\033[92mSuccessful login to WandB\033[00m")
+
+    training_pipeline(args)
+    print("\033[92mFinish training pipeline\033[00m")
+
+if __name__=='__main__':
+    main()

setup.sh

@@ -0,0 +1,8 @@
+echo "Hello"
+echo "..."
+pip install -q --upgrade pip
+pip install -q -U datasets
+pip install -q transformers
+pip install -q -r "/kaggle/working/dialogue-text-summarization/requirements.txt"
+echo "---------"
+echo "Set up complete!"

src/data/create_dataset.py

@@ -0,0 +1,115 @@
+
+import sys, os
+
+import argparse
+
+from bert_score import BERTScorer
+
+from transformers import (
+    T5Tokenizer, 
+    T5ForConditionalGeneration,
+    AutoTokenizer
+)
+
+import warnings
+warnings.filterwarnings("ignore")
+
+from huggingface_hub import login
+
+from datasets import load_dataset, Dataset
+
+path = os.path.abspath(os.path.dirname(__file__))
+sys.path.insert(0, path)
+
+from preprocessing import *
+
+def create_qds_triplet(datapath, split, start_index, end_index) -> Dataset:
+    data = load_dataset(datapath, split=split)
+    data = Dataset.from_dict(data[start_index:end_index])
+
+    scorer = BERTScorer(lang="en", rescale_with_baseline=True)
+
+    CHECKPOINT = "google/flan-t5-large"
+    tokenizer = T5Tokenizer.from_pretrained(CHECKPOINT)
+    model = T5ForConditionalGeneration.from_pretrained(CHECKPOINT)
+
+    qds_triplet = {
+        "query": [],
+        "dialogue": [],
+        "summary": []
+    }
+
+    dsp = DialogSumDataset(
+        tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
+    )
+
+    for dialogue, summary in zip(data["dialogue"], data["summary"]):
+        answerable_queries = []
+
+        while len(answerable_queries) < 1:
+            queries = dsp.generate_queries(model, tokenizer, summary, num_queries=5)
+
+            for query in queries:
+                ## Text based filtering
+                output = dsp.text_based_filtering(model, tokenizer, query, summary)
+                if "yes" in output.lower():
+                    answerable_queries.append(query)
+
+        n = len(answerable_queries)
+        print("Length of answerable queries:", n, end="  ###  ")
+
+        if n == 1:
+            qds_triplet["query"].append(answerable_queries[0])
+            qds_triplet["dialogue"].append(dialogue)
+            qds_triplet["summary"].append(summary)
+
+        if n > 1:
+            filtered_queries = []
+            scores = [[0.0]*n for _ in range(n)]
+
+            for i in range(n):
+                for j in range(n):
+                    if i > j:
+                        scores[i][j] = dsp.semantic_filtering(scorer, answerable_queries[i], answerable_queries[j])
+            
+            keep_indices = set(range(n))
+            for i in range(n):
+                for j in range(n):
+                    if scores[i][j] > 0.7 and i > j:
+                        keep_indices.discard(j)
+            
+            for i in sorted(keep_indices):
+                filtered_queries.append(answerable_queries[i])
+            
+            print("Length of filtered queries:", len(filtered_queries), end="  ###  ")
+
+            for query in filtered_queries:
+                qds_triplet["query"].append(query)
+                qds_triplet["dialogue"].append(dialogue)
+                qds_triplet["summary"].append(summary)
+
+        print("Length of inputs:", len(qds_triplet["summary"]))
+
+    return Dataset.from_dict(qds_triplet)
+
+if __name__=="__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--datapath", type=str, default="knkarthick/dialogsum")
+    parser.add_argument("--huggingface_hub_token", type=str, default="")
+    parser.add_argument("--split", type=str, default="train")
+    parser.add_argument("--start_index", type=int, default=0)
+    parser.add_argument("--end_index", type=int, default=-1)
+    args = parser.parse_args()
+
+    print("=========================================")
+    print('\n'.join(f' + {k}={v}' for k, v in vars(args).items()))
+    print("=========================================")
+
+    login(token=args.huggingface_hub_token)
+    print("Successfully logged in to Huggingface Hub")
+
+    qds_triplet = create_qds_triplet(args.datapath, args.split, args.start_index, args.end_index)
+
+    save_name = f"dialogsum-{args.split}-{args.start_index}-{args.end_index}"
+    qds_triplet.push_to_hub(save_name)
+    print(f"Saved to: {save_name}")

src/data/ingest_data.py

@@ -0,0 +1,113 @@
+
+from datasets import load_dataset
+from datasets import DatasetDict, Dataset
+import random
+from transformers import set_seed
+
+
+def ingest_data(datapath: str) -> DatasetDict:
+    set_seed(42)
+
+    QDS_LIMIT = 6000
+    if "," in datapath:
+        datapaths = datapath.split(",")
+    
+    datapath1 = "binwang/InstructDS_datasets"
+    datapath2 = "binwang/InstructDS_datasets"
+
+    all_train_data = []
+    origin_train_dialogsum = load_dataset(datapath1, "DialogSum", split="train")
+    qds_dialogsum = load_dataset(datapath2, "DialogSum_QDS", split="train")
+
+    new_data1 = []
+    for sample in origin_train_dialogsum:
+        new_sample = {
+            "instruction": "Please summarize the following dialogue.",
+            "input": sample["dialogue"],
+            "output": sample["summary"]
+        }
+        new_data1.append(new_sample)
+    origin_train_dialogsum = new_data1
+    all_train_data.extend(origin_train_dialogsum)
+
+    print("Len of origin_train_dialogsum: ", len(origin_train_dialogsum))
+    print("Len of all train data 1: ", len(all_train_data))
+    
+    new_data2 = []
+    for sample in qds_dialogsum:
+        new_sample = {
+            "instruction": "Please answer the following question.",
+            "input": sample["dialogue"],
+            "output": sample["summary"]
+        }
+        new_data2.append(new_sample)
+    qds_dialogsum = new_data2
+    qds_dialogsum = random.sample(qds_dialogsum, QDS_LIMIT)
+    all_train_data.extend(qds_dialogsum)
+    print("Len of all train data 2: ", len(all_train_data))
+
+
+    naive_all_train_data_dict = {
+        "instruction": [item["instruction"] for item in all_train_data],
+        "input": [item["input"] for item in all_train_data],
+        "output": [item["output"] for item in all_train_data]
+    }
+
+    print("Len of naive_all_train_data_dict: ", len(naive_all_train_data_dict["instruction"]))
+
+    subset_train_data = all_train_data
+    with_len_train_data_dict = {
+        "instruction": [item["instruction"] + f" The output should be {len(item['output'].split())} words long." for item in subset_train_data],
+        "input": [item["input"] for item in subset_train_data],
+        "output": [item["output"] for item in subset_train_data]
+    }
+
+    print("Len of with_len_train_data_dict: ", len(with_len_train_data_dict["instruction"]))
+
+    all_train_data_dict = {
+        "instruction": naive_all_train_data_dict["instruction"] + with_len_train_data_dict["instruction"],
+        "input": naive_all_train_data_dict["input"] + with_len_train_data_dict["input"],
+        "output": naive_all_train_data_dict["output"] + with_len_train_data_dict["output"]
+    }
+
+    print("Len of all_train_data_dict: ", len(all_train_data_dict["instruction"]))
+
+    raw_train_data = Dataset.from_dict(all_train_data_dict)
+    train_data = raw_train_data.shuffle()
+
+    print(type(train_data))
+    print(train_data["instruction"][:10])
+    print(train_data["input"][:10])
+    print(train_data["output"][:10])
+
+    print("===================", len(train_data), "===================")
+
+    # Validation data
+    all_validation_data = []
+    origin_validation_dialogsum = load_dataset(datapath1, "DialogSum", split="validation")
+
+    new_data1 = []
+    for sample in origin_validation_dialogsum:
+        new_sample = {
+            "instruction": "Please summarize the following dialogue.",
+            "input": sample["dialogue"],
+            "output": sample["summary"]
+        }
+        new_data1.append(new_sample)
+    
+    origin_validation_dialogsum = new_data1
+    all_validation_data.extend(origin_validation_dialogsum)
+
+    all_validation_data_dict = {
+        "instruction": [item["instruction"] for item in all_validation_data],
+        "input": [item["input"] for item in all_validation_data],
+        "output": [item["output"] for item in all_validation_data]
+    }
+
+    raw_validation_data = Dataset.from_dict(all_validation_data_dict)
+    validation_data = raw_validation_data.shuffle()
+
+    return DatasetDict({
+        "train": train_data,
+        "validation": validation_data
+    })

src/data/merge_dataset.py

@@ -0,0 +1,41 @@
+
+import os, sys
+
+import argparse
+
+from datasets import load_dataset, concatenate_datasets, Dataset
+from huggingface_hub import login
+
+path = os.path.abspath(os.path.dirname(__file__))
+sys.path.insert(0, path)
+
+def merge_dataset(datapaths) -> Dataset:
+    datapaths = datapaths.split(",")
+    dataset = load_dataset(datapaths[0], split="train")
+
+    for i in range(1, len(datapaths)):
+        data = load_dataset(datapaths[i], split="train")
+        data = concatenate_datasets([dataset, data])
+    
+    return dataset
+
+
+if __name__=="__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--datapaths", type=str, default="")
+    parser.add_argument("--huggingface_hub_token", type=str, default="")
+    parser.add_argument("--split", type=str, default="train")
+    args = parser.parse_args()
+
+    print("=========================================")
+    print('\n'.join(f' + {k}={v}' for k, v in vars(args).items()))
+    print("=========================================")
+
+    login(token=args.huggingface_hub_token)
+    print("Successfully logged in to Huggingface Hub")
+
+    dataset = merge_dataset(datapaths=args.datapaths)
+    
+    DATASET_ID = "qds-triplet-dialogsum"
+    dataset.push_to_hub(DATASET_ID)
+    print(f"Successful push to Huggingface Hub: {DATASET_ID}")

src/data/preprocessing.py

@@ -0,0 +1,113 @@
+
+from datasets import DatasetDict, Dataset
+import random
+from bert_score import BERTScorer
+
+from transformers import (
+    T5Tokenizer,
+    T5ForConditionalGeneration
+)
+
+class DialogSumDataset:
+    def __init__(self, tokenizer, use_contrastive_loss=False, tokenizing_strategy=1) -> None:
+        self.tokenizer = tokenizer
+        self.use_contrastive_loss = use_contrastive_loss
+        self.tokenizing_strategy = tokenizing_strategy
+        
+    def handle_data(self, data: DatasetDict) -> DatasetDict:
+        try:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            tokenized_dataset = data.map(self.preprocess_function, batched=True)
+            tokenized_dataset = tokenized_dataset.remove_columns([key for key in data["train"][0].keys()])
+
+            print("+++++++++++++++++++")
+            print(tokenized_dataset)
+            print("+++++++++++++++++++")
+            
+            return tokenized_dataset
+
+        except Exception as e:
+            print(f"\033[31m\nError while tokenizing data: {e}\033[00m")
+            raise e
+
+    def preprocess_function(self, data: Dataset) -> Dataset:
+        ###
+        if self.tokenizing_strategy<=2:
+            prefix = "Summarize the following conversation:\n###\n"
+            suffix = "\n###\nSummary: "
+            inputs = [prefix + input + suffix for input in data["dialogue"]]
+            targets = data["summary"]
+            
+            if self.tokenizing_strategy==1:
+                max_source_length = 1024
+                max_target_length = 176
+
+            if self.tokenizing_strategy==2:
+                max_source_length = 1224
+                max_target_length = 176
+
+        if self.tokenizing_strategy==3:
+            inputs = ["### Instruction: " + instruction + "\n### Input: " + input + "\n### Response: " for instruction, input in zip(data["instruction"], data["input"])]
+            targets = data["output"]
+
+            max_source_length = 1024
+            max_target_length = 176
+
+        data["input_ids"] = self.tokenizer(inputs, max_length=max_source_length, padding="max_length", truncation=True, return_tensors="pt").input_ids
+        # data["attention_mask"] = self.tokenizer(inputs, max_length=max_source_length, padding="max_length", truncation=True, return_tensors="pt").attention_mask
+        data["labels"] = self.tokenizer(targets, max_length=max_target_length, padding="max_length", truncation=True, return_tensors="pt").input_ids
+        
+        # Generate negative examples:
+        if self.use_contrastive_loss==True:
+            negative_summaries = self.generate_negative_examples(data["summary"])
+            data["negative_labels"] = self.tokenizer(negative_summaries, max_length=max_target_length, padding="max_length", truncation=True, return_tensors="pt").input_ids
+            print("Complete generate negative examples!")
+
+        label_ignore_ids = []
+        for label in data["labels"]:
+            label_example = [l if l != 0 else -100 for l in label]
+            label_ignore_ids.append(label_example)
+
+        data["labels"] = label_ignore_ids
+
+        return data
+    
+    ## Create Negetive Example for Contrastive Learning
+    def generate_negative_examples(self, summaries):
+        negative_summaries = []
+        for summary in summaries:
+            words = summary.split()
+            random.shuffle(words)
+            negative_summaries.append(" ".join(words))
+        return negative_summaries
+
+    ## Create Instruction Dataset
+    def generate_queries(self, model, tokenizer, summary, num_queries):
+        input_text = "Generate an answerable and specific question based on the following context:. ###\nContext: " + summary
+        input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+        outputs = model.generate(input_ids, max_length=64, num_return_sequences=num_queries, do_sample=True)
+        queries = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
+        return queries
+    
+    def text_based_filtering(self, model, tokenizer, query, summary):
+        input_text = "Is the question fully answerable from the context without any guessing, yes or no?###\nQuestion: " + query + "###\nContext: " + summary + "###Answer: "
+        input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+        output_ids = model.generate(input_ids, num_return_sequences=1)
+        output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        return output_text
+    
+    def semantic_filtering(self, scorer, query1, query2):
+        score = scorer.score([query1], [query2])[0]
+        return score
+
+
+def preprocessing_data(data: DatasetDict, tokenizer, use_contrastive_loss=False, tokenizing_strategy=False) -> DatasetDict:
+    try:
+        dataset_ds = DialogSumDataset(tokenizer, use_contrastive_loss, tokenizing_strategy)
+        tokenized_data = dataset_ds.handle_data(data)
+
+        return tokenized_data
+
+    except Exception as e:
+        print(f"\nError while pre-processing data: {e}")
+        raise e

src/evaluate/evaluation.py

@@ -0,0 +1,81 @@
+import os
+import sys
+
+from datasets import Dataset
+
+import evaluate
+import torch
+
+import logging
+
+# =  =  =  =  =  =  =  =  =  =  =  Logging Setup  =  =  =  =  =  =  =  =  =  =  =  =  = 
+logger = logging.getLogger(__name__)
+logging.basicConfig(
+    format  = "%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt = "%m/%d/%Y %H:%M:%S",
+    level   = logging.INFO,
+)
+# =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  =  = 
+
+from transformers import AutoModelForSeq2SeqLM
+
+path = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.insert(0, path)
+
+from model.model import Model
+
+
+class RougeEvaluation:
+    def __init__(self) -> None:
+        self.rouge_metric = evaluate.load("rouge")
+        
+    def compute_rouge_metric(self, generated_summary, reference_summary) -> dict:
+        results = self.rouge_metric.compute(
+            predictions=generated_summary,
+            references=reference_summary,
+            use_aggregator=True,
+            use_stemmer=True
+        )
+        return results
+    
+
+def evaluation_rouge(model: Model, data: Dataset, generation_config) -> dict:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.base_model = model.get_model()
+
+    dialogues = data["dialogue"]
+
+    human_summaries = [summary for summary in data["summary"]]
+
+    model_summaries = []
+
+    prefix = "Summarize the following dialogue:\n###\n"
+    suffix = "\n### Summary: "
+
+    # print("\n******************************")
+    # idx = 0
+    # for answer, dialogue in zip(data["answer"], data["dialogue"]):
+    #     prefix = "Please summarize the following dialogue focused on the context query:"
+    #     input = prefix + "\n### Queryr: " + answer + "\n### Dialogue: " + dialogue + "\n### The summary should be around " + str(int(0.2*len(dialogue.split()))) + " words." + "\n### Summary: "
+
+    for idx, dialogue in enumerate(dialogues):
+        input = prefix + dialogue + suffix
+        
+        print(idx, end="# ")
+        output_text = model.generate_summary(input, generation_config, do_sample=False)
+
+        model_summaries.append(output_text)
+        idx += 1
+
+    logger.info("Evaluating summaries...")
+
+    rouge_evaluator = RougeEvaluation()
+
+    results = rouge_evaluator.compute_rouge_metric(model_summaries, human_summaries)
+
+    generated_lengths = [len(summary.split()) for summary in model_summaries]
+    average_gen_len = sum(generated_lengths) / len(generated_lengths) if generated_lengths else 0
+
+    results["gen_len"] = average_gen_len
+    
+    return results

src/evaluate/rouge_metric.py

@@ -0,0 +1,52 @@
+import evaluate
+import nltk
+import numpy as np
+from nltk.tokenize import sent_tokenize
+
+from transformers import AutoTokenizer
+
+import os
+import sys
+
+path = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.insert(0, path)
+
+
+def postprocess_text(preds, labels):
+    nltk.download("punkt")
+
+    preds = [pred.strip() for pred in preds]
+    labels = [label.strip() for label in labels]
+
+    preds = ["\n".join(sent_tokenize(pred)) for pred in preds]
+    labels = ["\n".join(sent_tokenize(label)) for label in labels]
+
+    return preds, labels
+
+
+def compute_metrics(eval_preds, tokenizer, metric):
+    preds, labels = eval_preds
+    if isinstance(preds, tuple):
+        preds = preds[0]
+
+    decoded_preds = tokenizer.batch_decode(preds, 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)
+
+    # Some simple post-processing
+    decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
+
+    # metric = evaluate.load("rouge")
+    rouge_results = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
+    rouge_results = {k: round(v * 100, 4) for k, v in rouge_results.items()}
+    
+    results = {
+        "rouge1": rouge_results["rouge1"],
+        "rouge2": rouge_results["rouge2"],
+        "rougeL": rouge_results["rougeL"],
+        "rougeLsum": rouge_results["rougeLsum"],
+        "gen_len": np.mean([np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds])
+    }
+
+    return results

src/model/model.py

@@ -0,0 +1,80 @@
+import torch
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+)
+
+from peft import (
+    get_peft_model,
+)
+
+class Model:
+    def __init__(self, checkpoint):
+        self.checkpoint = checkpoint
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.checkpoint)
+        self.base_model = None
+
+    def get_model(self):
+        return AutoModelForSeq2SeqLM.from_pretrained(self.checkpoint)
+
+    def get_peft(self, lora_config):
+        return get_peft_model(self.base_model, lora_config)
+    
+    def prepare_quantize(self, bnb_config):
+        return AutoModelForSeq2SeqLM.from_pretrained(self.checkpoint, 
+                                                                 quantization_config=bnb_config, 
+                                                                 device_map={"":0}, 
+                                                                 trust_remote_code=True)
+        # self.base_model.gradient_checkpointing_enable()
+        # self.base_model = prepare_model_for_kbit_training(self.base_model)
+
+
+    def generate_summary(self, input_text, generation_config, do_sample=True):
+        input_ids = self.tokenizer.encode(input_text, return_tensors="pt", max_length=1024, truncation=True, padding="max_length")
+        output_ids = self.base_model.generate(input_ids=input_ids, do_sample=do_sample, generation_config=generation_config)
+        
+        if "bart" in self.checkpoint:
+            output_ids[0][1] = 2
+
+        output_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        print(f"\033[94mSummary: {output_text}\n\033[00m")
+        return output_text
+
+class BartSum(Model):
+    def __init__(self, checkpoint):
+        super().__init__(checkpoint)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.checkpoint)
+
+    def get_model(self):
+        return AutoModelForSeq2SeqLM.from_pretrained(self.checkpoint)
+
+
+class FlanT5Sum(Model):
+    def __init__(self, checkpoint):
+        super().__init__(checkpoint)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.checkpoint)
+
+    def get_model(self):
+        return AutoModelForSeq2SeqLM.from_pretrained(self.checkpoint)
+
+
+def load_model(checkpoint):
+
+    try:
+        if "bart" in checkpoint:
+            print(f"\033[92mLoad Bart model from checkpoint: {checkpoint}\033[00m")
+            return BartSum(checkpoint)
+        
+        if "flan" in checkpoint:
+            print(f"\033[92mLoad Flan-T5 model from checkpoint: {checkpoint}\033[00m")
+            return FlanT5Sum(checkpoint)
+        
+        else:
+            print(f"\033[92mLoad general model from checkpoint: {checkpoint}\033[00m")
+            return Model(checkpoint)
+        
+    except Exception as e:
+        print("Error while loading model: {e}")
+        raise e

src/pipelines/training_pipeline.py

@@ -0,0 +1,168 @@
+import os
+import sys
+import argparse
+import numpy as np
+import nltk
+
+from nltk.tokenize import sent_tokenize
+from transformers import (
+    Seq2SeqTrainer,
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM
+)
+
+from peft import get_peft_model, prepare_model_for_kbit_training
+
+path = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.insert(0, path)
+
+from utils import *
+
+# from model.models import load_model
+from model.model import load_model
+from data.preprocessing import preprocessing_data
+from data.ingest_data import ingest_data
+
+import evaluate
+
+
+def training_pipeline(args: argparse.Namespace):
+    try:
+        print("=========================================")
+        print('\n'.join(f' + {k}={v}' for k, v in vars(args).items()))
+        print("=========================================")
+
+        import torch
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        model = load_model(args.checkpoint)
+        tokenizer = AutoTokenizer.from_pretrained(args.checkpoint)
+        print(type(tokenizer))
+        
+        if (args.lora == False):
+            print("lora=Fasle, quantize=False")
+            base_model = AutoModelForSeq2SeqLM.from_pretrained(args.checkpoint).to(device)
+            # model.base_model = model.get_model()
+            # model.base_model.to(device)
+
+        else:
+            from peft import LoraConfig, TaskType
+            from transformers import BitsAndBytesConfig
+            import torch
+            # Define LoRA Config 
+            lora_config = LoraConfig(
+                r=args.lora_rank, 
+                lora_alpha=args.lora_alpha,
+                target_modules=args.target_modules.split(","),
+                lora_dropout=args.lora_dropout,
+                bias="none",
+                task_type=TaskType.SEQ_2_SEQ_LM
+            )
+
+            if (args.quantize == True):
+                print("Quantize=True, lora=True")
+                bnb_config = BitsAndBytesConfig(
+                    load_in_4bit=True,
+                    bnb_4bit_use_double_quant=True,
+                    bnb_4bit_quant_type="nf4",
+                    bnb_4bit_compute_dtype=torch.bfloat16
+                )
+
+                base_model = AutoModelForSeq2SeqLM.from_pretrained(args.checkpoint,
+                                                                   quantization_config=bnb_config,
+                                                                   device_map={"":0},
+                                                                   trust_remote_code=True)
+                base_model = prepare_model_for_kbit_training(base_model)
+
+            if (args.quantize==False):
+                print("Quantize=False, lora=True")
+                base_model = AutoModelForSeq2SeqLM.from_pretrained(args.checkpoint).to(device)
+
+            # add LoRA adaptor
+            print("Base model:", model.base_model)
+            base_model = get_peft_model(base_model, lora_config)
+            base_model.print_trainable_parameters()
+
+
+        # Load data from datapath
+        data = ingest_data(args.datapath)
+        print("\033[92m[+] Complete loading dataset!\033[00m")
+
+        # Pre-processing data
+        data = preprocessing_data(data, tokenizer, use_contrastive_loss=args.use_contrastive_loss, tokenizing_strategy=args.tokenizing_strategy)
+        print("\033[92m[+] Complete pre-processing dataset!\033[00m")
+
+        # Load training arguments
+        training_args = load_training_arguments(args)
+        print("\033[92m[+] Complete loading training arguments!\033[00m")
+
+        # Load metric
+        metric = evaluate.load("rouge")
+        nltk.download("punkt")
+
+        def postprocess_text(preds, labels):
+            preds = [pred.strip() for pred in preds]
+            labels = [label.strip() for label in labels]
+
+            preds = ["\n".join(sent_tokenize(pred)) for pred in preds]
+            labels = ["\n".join(sent_tokenize(label)) for label in labels]
+
+            return preds, labels
+        
+        def compute_metric(eval_preds):
+            preds, labels = eval_preds
+            if isinstance(preds, tuple):
+                preds = preds[0]
+            
+            decoded_preds = tokenizer.batch_decode(preds, 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)
+
+            # Some simple post-processing
+            decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
+
+            # metric = evaluate.load("rouge")
+            rouge_results = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
+            rouge_results = {k: round(v * 100, 4) for k, v in rouge_results.items()}
+            
+            results = {
+                "rouge1": rouge_results["rouge1"],
+                "rouge2": rouge_results["rouge2"],
+                "rougeL": rouge_results["rougeL"],
+                "rougeLsum": rouge_results["rougeLsum"],
+                "gen_len": np.mean([np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds])
+            }
+
+            return results
+        
+        # Load trainer
+        if args.use_contrastive_loss==True:
+            trainer = ContrastiveLearningTrainer(model=base_model,
+                                     train_dataset=data["train"],
+                                     eval_dataset=data["validation"],
+                                     tokenizer=tokenizer,
+                                     compute_metrics=compute_metric)
+
+        if args.use_contrastive_loss==False:
+            trainer = Seq2SeqTrainer(model=base_model,
+                                args=training_args,
+                                train_dataset=data["train"],
+                                eval_dataset=data["validation"],
+                                tokenizer=tokenizer,
+                                compute_metrics=compute_metric)
+        
+        print("\033[92m[+] Complete loading trainer!\033[00m")
+
+        # Train model
+        trainer.train()
+        print("\033[92m[+] Complete training!\033[00m")
+
+        # Push to Huggingface Hub
+        trainer.push_to_hub()
+        print("\033[92m [+] Complete pushing model to hub!\033[00m")
+
+    except Exception as e:
+        print(f"\033[31m\nError while training: {e}\033[00m")
+        raise e
+    

src/utils.py

@@ -0,0 +1,131 @@
+import argparse
+
+import os
+import sys
+
+import torch
+import torch.nn as nn
+
+from transformers import (
+    Seq2SeqTrainingArguments, 
+    Seq2SeqTrainer,
+)
+
+path = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.insert(0, path)
+
+# from src.evaluate.rouge_metric import compute_metrics
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Fine tuning LLM for Dialogue Text Summarization")
+    parser.add_argument("--huggingface_hub_token", type=str, default=None)
+    parser.add_argument("--wandb_token", type=str, default=None)
+
+    parser.add_argument("--checkpoint", type=str, default="google/flan-t5-base")
+    parser.add_argument("--datapath", type=str, default="knkarthick/dialogsum")
+
+    parser.add_argument("--output_dir", type=str, default="fine-tuned-flant5")
+    parser.add_argument("--overwrite_output_dir", action="store_true")
+    
+    parser.add_argument("--num_train_epochs", type=int, default=3)
+    parser.add_argument("--per_device_train_batch_size", type=int, default=4)
+    parser.add_argument("--per_device_eval_batch_size", type=int, default=4)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=2)
+
+    parser.add_argument("--learning_rate", type=float, default=0.00005)
+    parser.add_argument("--weight_decay", type=float, default=0.005)
+
+    parser.add_argument("--evaluation_strategy", type=str, default="no")
+    parser.add_argument("--save_strategy", type=str, default="no")
+
+    parser.add_argument("--logging_strategy", type=str, default="steps")
+    parser.add_argument("--logging_steps", type=int, default=1000)
+    parser.add_argument("--save_total_limit", type=int, default=1)
+
+    parser.add_argument("--report_to", type=str, default="wandb")
+    parser.add_argument("--run_name", type=str, default="flan-t5-base-model")
+
+    parser.add_argument("--predict_with_generate", action="store_true")
+
+    parser.add_argument("--min_new_tokens", type=int, default=10)
+    parser.add_argument("--max_new_tokens", type=int, default=256)
+    parser.add_argument("--temperature", type=float, default=0.9)
+    parser.add_argument("--top_p", type=float, default=1.0)
+    parser.add_argument("--top_k", type=int, default=50)
+
+    parser.add_argument("--lora", action="store_true")
+    parser.add_argument("--quantize", action="store_true")
+
+    parser.add_argument("--lora_rank", type=int, default=8)
+    parser.add_argument("--lora_alpha", type=int, default=16)
+    parser.add_argument("--target_modules", type=str, default="q,v")
+    parser.add_argument("--lora_dropout", type=float, default=0.05)
+
+    parser.add_argument("--use_contrastive_loss", action="store_true")
+    parser.add_argument("--tokenizing_strategy", type=int, default=1)
+
+    args = parser.parse_args()
+    return args
+
+
+def load_training_arguments(args):
+    try:
+        training_args = Seq2SeqTrainingArguments(
+            output_dir=args.output_dir,
+            overwrite_output_dir=args.overwrite_output_dir,
+
+            num_train_epochs=args.num_train_epochs,
+            per_device_train_batch_size=args.per_device_train_batch_size,
+            per_device_eval_batch_size=args.per_device_eval_batch_size,
+            gradient_accumulation_steps=args.gradient_accumulation_steps,
+            
+            learning_rate=args.learning_rate,
+            weight_decay=args.weight_decay,
+
+            evaluation_strategy=args.evaluation_strategy,
+            save_strategy=args.save_strategy,
+            
+            logging_strategy=args.logging_strategy,
+            logging_steps=args.logging_steps,
+            save_total_limit=args.save_total_limit,
+            
+            report_to=args.report_to,
+            run_name=args.run_name,
+
+            predict_with_generate=args.predict_with_generate
+        )
+
+        return training_args
+    
+    except Exception as e:
+        print(f"Error while loading training arguments: {e}")
+        raise e
+
+class ContrastiveLoss(nn.Module):
+    def __init__(self, margin=1.0):
+        super(ContrastiveLoss, self).__init__()
+        self.margin = margin
+        self.cosine_similarity = nn.CosineSimilarity(dim=1, eps=1e-6)
+
+    def forward(self, dialgue_embeddings, pos_summary_embeddings, neg_summary_embeddings):
+        pos_sim = self.cosine_similarity(dialgue_embeddings, pos_summary_embeddings)
+        neg_sim = self.cosine_similarity(dialgue_embeddings, neg_summary_embeddings)
+        loss = torch.mean(1-pos_sim) + torch.clamp(neg_sim-self.margin, min=0.0)
+
+        return loss
+
+class ContrastiveLearningTrainer(Seq2SeqTrainer):
+    def compute_loss(model, inputs, return_outputs=False):
+        output = model(**inputs)
+        lm_loss = output.loss
+
+        dialogue_embeddings = model.encoder(inputs["input_ids"]).last_hidden_state
+        pos_summary_embeddings = model.encoder(inputs["labels"]).last_hidden_state
+        neg_summary_embeddings = model.encoder(inputs["negative_labels"]).last_hidden_state
+
+        contrastive_loss = ContrastiveLoss(margin=1.0)(dialogue_embeddings, pos_summary_embeddings, neg_summary_embeddings)
+
+        # Combine losses
+        total_loss = lm_loss + contrastive_loss
+
+        return (total_loss, output) if return_outputs else total_loss

test_streaming.py

@@ -0,0 +1,76 @@
+import streamlit as st
+import replicate
+import os
+from transformers import AutoTokenizer, GenerationConfig, AutoModelForSeq2SeqLM
+import torch
+
+# Set Replicate API token
+with st.sidebar:
+    st.title('Dialogue Text Summarization')
+    if 'REPLICATE_API_TOKEN' in st.secrets:
+        replicate_api = st.secrets['REPLICATE_API_TOKEN']
+    else:
+        replicate_api = st.text_input('Enter Replicate API token:', type='password')
+        if not (replicate_api.startswith('r8_') and len(replicate_api) == 40):
+            st.warning('Please enter your Replicate API token.', icon='⚠️')
+            st.markdown("**Don't have an API token?** Head over to [Replicate](https://replicate.com) to sign up for one.")
+
+    os.environ['REPLICATE_API_TOKEN'] = replicate_api
+    st.subheader("Adjust model parameters")
+    min_new_tokens = st.slider('Min new tokens', min_value=1, max_value=256, step=1, value=10)
+    temperature = st.slider('Temperature', min_value=0.01, max_value=1.00, step=0.01, value=1.0)
+    top_k = st.slider('Top_k', min_value=1, max_value=50, step=1, value=20)
+    top_p = st.slider('Top_p', min_value=0.01, max_value=1.00, step=0.01, value=1.0)
+
+# Initialize model and tokenizer
+checkpoint = "dtruong46me/train-bart-base"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint).to(device)
+
+st.title("Dialogue Text Summarization")
+st.caption("Natural Language Processing Project 20232")
+st.write("---")
+
+input_text = st.text_area("Dialogue", height=200)
+
+generation_config = GenerationConfig(
+    min_new_tokens=min_new_tokens,
+    max_new_tokens=320,
+    temperature=temperature,
+    top_p=top_p,
+    top_k=top_k
+)
+
+def generate_summary(model, input_text, generation_config, tokenizer):
+    prefix = "Summarize the following conversation: \n\n###"
+    suffix = "\n\nSummary:"
+    input_ids = tokenizer.encode(prefix + input_text + suffix, return_tensors="pt").to(model.device)
+    prompt_str = tokenizer.decode(input_ids[0], skip_special_tokens=True)
+    return prompt_str
+
+def stream_summary(prompt_str, temperature, top_p):
+    for event in replicate.stream(
+        "snowflake/snowflake-arctic-instruct",
+        input={"prompt": prompt_str,
+               "prompt_template": r"{prompt}",
+               "temperature": temperature,
+               "top_p": top_p}):
+        yield str(event['output'])
+
+if st.button("Submit"):
+    st.write("---")
+    st.write("## Summary")
+
+    if not replicate_api:
+        st.error("Please enter your Replicate API token!")
+    elif not input_text:
+        st.error("Please enter a dialogue!")
+    else:
+        prompt_str = generate_summary(model, input_text, generation_config, tokenizer)
+        summary_container = st.empty()
+
+        summary_text = ""
+        for output in stream_summary(prompt_str, temperature, top_p):
+            summary_text += output
+            summary_container.text(summary_text)