Spaces:
Runtime error
Runtime error
| # -*- coding: utf-8 -*- | |
| """ | |
| Created on Thu May 19 13:22:32 2022 | |
| @author: UTKARSH | |
| """ | |
| import glob | |
| import os | |
| import pandas as pd | |
| import torch | |
| from torch.utils.data import Dataset, DataLoader | |
| from sklearn.model_selection import train_test_split | |
| from src.clean import clean_license_text | |
| from tqdm.auto import tqdm | |
| from transformers import ( | |
| AdamW, | |
| T5ForConditionalGeneration, | |
| T5TokenizerFast as T5Tokenizer, | |
| AutoModelForSeq2SeqLM, | |
| AutoTokenizer | |
| ) | |
| from src.read_data import read_license_summary_data | |
| MODEL_PATH = "models/" | |
| MODEL_FILENAME = "t5-base.model" | |
| MODEL_NAME = "t5-base" | |
| TOKENIZER = None | |
| TEXT_MAX_TOKEN_LEN = 512 | |
| SUMMARY_MAX_TOKEN_LEN = 128 | |
| N_EPOCHS = 1 | |
| BATCH_SIZE = 1 | |
| device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") | |
| class LicenseSummaryDataset(Dataset): | |
| def __init__( | |
| self, | |
| data: pd.DataFrame, | |
| tokenizer: T5Tokenizer, | |
| text_max_token_len: int=512, | |
| summary_max_token_len: int=128 | |
| ): | |
| self.tokenizer = tokenizer | |
| self.data = data | |
| self.text_max_token_len = text_max_token_len | |
| self.summary_max_token_len = summary_max_token_len | |
| def __len__(self): | |
| return len(self.data) | |
| def __getitem__(self, index: int): | |
| data_row = self.data.iloc[index] | |
| text = data_row["text"] | |
| text_encoding = self.tokenizer( | |
| text, | |
| max_length=self.text_max_token_len, | |
| padding="max_length", | |
| truncation=True, | |
| return_attention_mask=True, | |
| add_special_tokens=True, | |
| return_tensors="pt" | |
| ) | |
| summary_encoding = self.tokenizer( | |
| data_row["summary"], | |
| max_length=self.summary_max_token_len, | |
| padding="max_length", | |
| truncation=True, | |
| return_attention_mask=True, | |
| add_special_tokens=True, | |
| return_tensors="pt" | |
| ) | |
| labels = summary_encoding["input_ids"] | |
| labels[labels == 0] = -100 | |
| return dict( | |
| text=text, | |
| summary=data_row["summary"], | |
| text_input_ids=text_encoding["input_ids"].flatten(), | |
| text_attention_mask=text_encoding["attention_mask"].flatten(), | |
| labels=labels.flatten(), | |
| labels_attention_mask=summary_encoding["attention_mask"].flatten() | |
| ) | |
| def prepare_dataloaders(): | |
| """ | |
| Helper method to load data and create batched Dataloaders | |
| Returns | |
| ------- | |
| train_dataloader : DataLoader | |
| Train DataLoader. | |
| dev_dataloader : DataLoader | |
| Validation DataLoader. | |
| """ | |
| license_summary_data = pd.DataFrame(read_license_summary_data()) | |
| train_df, dev_df = train_test_split(license_summary_data, test_size=0.1) | |
| TOKENIZER = T5Tokenizer.from_pretrained(MODEL_NAME) | |
| train_dataset = LicenseSummaryDataset( | |
| train_df, | |
| TOKENIZER, | |
| TEXT_MAX_TOKEN_LEN, | |
| SUMMARY_MAX_TOKEN_LEN | |
| ) | |
| dev_dataset = LicenseSummaryDataset( | |
| dev_df, | |
| TOKENIZER, | |
| TEXT_MAX_TOKEN_LEN, | |
| SUMMARY_MAX_TOKEN_LEN | |
| ) | |
| train_dataloader = DataLoader( | |
| train_dataset, | |
| batch_size=BATCH_SIZE, | |
| shuffle=True, | |
| num_workers=0 | |
| ) | |
| dev_dataloader = DataLoader( | |
| dev_dataset, | |
| batch_size=BATCH_SIZE, | |
| shuffle=True, | |
| num_workers=0 | |
| ) | |
| return train_dataloader, dev_dataloader | |
| def train(epoch, model, dataloader, optimizer, batch_size): | |
| """ | |
| Trains the given model on the given data for the given number of epochs. | |
| Parameters | |
| ---------- | |
| epoch : int | |
| The epoch number for which the model is being trained. | |
| model : Summarizer | |
| A summarizer model which we train. | |
| dataloader : torch.utils.data.DataLoader | |
| The dataloader on which the model is to be trained. | |
| optimizer : transformers.AdamW | |
| The optimizer to be used to optimize weights during training. | |
| batch_size : int | |
| The size of each batch as set in the dataloader. | |
| """ | |
| model.train() | |
| total_train_loss = 0 | |
| for _, batch in tqdm(enumerate(dataloader)): | |
| model.zero_grad() | |
| print(_) | |
| input_ids = batch["text_input_ids"].to(device, dtype=torch.long) | |
| attention_mask = batch["text_attention_mask"].to(device, dtype=torch.long) | |
| labels = batch["labels"].to(device, dtype=torch.long) | |
| labels_attention_mask = batch["labels_attention_mask"].to(device, dtype=torch.long) | |
| model_output = model( | |
| input_ids=input_ids, | |
| attention_mask=attention_mask, | |
| decoder_attention_mask=labels_attention_mask, | |
| labels=labels | |
| ) | |
| # loss, _ = model_output.loss, model_output.logits | |
| loss = model_output.loss | |
| total_train_loss += loss.item() | |
| loss.backward() | |
| # torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) | |
| optimizer.step() | |
| avg_train_loss = total_train_loss / len(dataloader) | |
| print(f"Epoch {epoch}: Training loss: {avg_train_loss}") | |
| def train_and_save_model(train_dataloader, PATH): | |
| """ | |
| Trains a summarizer model from the given Dataloader and saves it at the | |
| given path | |
| Parameters | |
| ---------- | |
| train_dataloader : Dataloader | |
| Batched Training Dataloader. | |
| PATH : str | |
| Path where the trained model is to be saved. | |
| Returns | |
| ------- | |
| model : Summarizer / torch.nn.Module | |
| Trained model. | |
| """ | |
| model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME, return_dict=True).to(device) | |
| optimizer = AdamW(model.parameters(), lr=3e-5) | |
| for epoch in range(1, N_EPOCHS + 1): | |
| train(epoch, model, train_dataloader, optimizer, BATCH_SIZE) | |
| torch.save(model.state_dict(), PATH) | |
| print("Model Saved!") | |
| return model | |
| def summarize_text_with_model(text, model, tokenizer): | |
| """ | |
| Summarizes License text using the given trained T5 model. | |
| Parameters | |
| ---------- | |
| text : str | |
| The License text to be summarized. | |
| model : Summarizer / torch.nn.Module | |
| The trained model which is to be used to summarize text. | |
| tokenizer : Tokenzier | |
| The tokenizer used to tokenize text for model. | |
| Returns | |
| ------- | |
| str | |
| Summary of the License text from the given model. | |
| definitions : str | |
| Definitions extracted from the License text. | |
| """ | |
| text, definitions, _ = clean_license_text(text) | |
| text_encoding = tokenizer( | |
| text, | |
| max_length=TEXT_MAX_TOKEN_LEN, | |
| padding="max_length", | |
| truncation=True, | |
| return_attention_mask=True, | |
| add_special_tokens=True, | |
| return_tensors="pt" | |
| ) | |
| generated_ids = model.generate( | |
| input_ids=text_encoding["input_ids"].to(device, dtype=torch.long), | |
| attention_mask=text_encoding["attention_mask"].to(device, dtype=torch.long), | |
| max_length=SUMMARY_MAX_TOKEN_LEN, | |
| num_beams=2, | |
| repetition_penalty=2.5, | |
| length_penalty=1.0, | |
| early_stopping=True | |
| ) | |
| preds = [ | |
| tokenizer.decode( | |
| gen_id, | |
| skip_special_tokens=True, | |
| clean_up_tokenization_spaces=True | |
| ) for gen_id in generated_ids | |
| ] | |
| return "".join(preds), definitions | |
| def summarize(text, load_from_huggingface=True): | |
| """ | |
| Summarizes the given License text | |
| Parameters | |
| ---------- | |
| text : str | |
| Preprocessed License text. | |
| load_from_huggingface : boolean | |
| Toggles whether or not to load the model from huggingface. If set to | |
| False, this will load or train the model locally. | |
| Returns | |
| ------- | |
| summary : str | |
| Summary of the License text. | |
| definitions : str | |
| Definitions extracted from the License text. | |
| """ | |
| if load_from_huggingface: | |
| print("Loading Model from HuggingFace...") | |
| CUSTOM_MODEL_NAME = "utkarshsaboo45/ClearlyDefinedLicenseSummarizer" | |
| model = AutoModelForSeq2SeqLM.from_pretrained(CUSTOM_MODEL_NAME).to(device) | |
| tokenizer = AutoTokenizer.from_pretrained(CUSTOM_MODEL_NAME) | |
| else: | |
| if os.path.exists(MODEL_PATH + MODEL_FILENAME): | |
| print("Loading Model...") | |
| model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME, return_dict=True).to(device) | |
| TOKENIZER = T5Tokenizer.from_pretrained(MODEL_NAME) | |
| model.load_state_dict(torch.load(MODEL_PATH + MODEL_FILENAME)) | |
| model.eval() | |
| else: | |
| print("Training model...") | |
| if not os.path.exists(MODEL_PATH): | |
| os.makedirs(MODEL_PATH) | |
| train_dataloader, _ = prepare_dataloaders() | |
| model = train_and_save_model(train_dataloader, MODEL_PATH + MODEL_FILENAME) | |
| tokenizer = TOKENIZER | |
| summary, definitions = summarize_text_with_model(text, model, tokenizer) | |
| return summary, definitions | |
| def summarize_license_files(path): | |
| """ | |
| Summarize License files from paths and save them as summary text files. | |
| Parameters | |
| ---------- | |
| path : list(str) | |
| A list of paths of the License files. | |
| """ | |
| paths = glob.glob(path + "*.txt") | |
| for license_path in paths: | |
| with open(license_path, "r", encoding="utf-8") as f: | |
| summary, _ = summarize(f.read()) | |
| with open(license_path.replace(".txt", "") + "__summary.txt", "w", encoding="utf-8") as f: | |
| f.write(summary) | |