Upload streamlit_app.py

streamlit_app.py

@@ -0,0 +1,255 @@
+import streamlit as st
+import numpy as np
+import pandas as pd
+
+# modeling
+import torch
+from torch.utils.data import Dataset, DataLoader
+import pytorch_lightning as pl
+from pytorch_lightning import Trainer, seed_everything
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+from transformers import (
+    T5ForConditionalGeneration,
+    T5TokenizerFast as T5Tokenizer,
+    )
+from transformers.optimization import Adafactor
+
+# aesthetics
+from IPython.display import Markdown, display, clear_output
+import re
+import warnings
+warnings.filterwarnings(
+    "ignore", ".*Trying to infer the `batch_size` from an ambiguous collection.*"
+)
+seed_everything(25429)
+
+# scoring
+import spacy
+from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(device)
+
+# loading the model
+hug = 't5-small'
+t5tokenizer = T5Tokenizer.from_pretrained(hug)
+t5model = T5ForConditionalGeneration.from_pretrained(hug, return_dict=True)
+
+# defining tokens
+SEP_TOKEN = '<sep>'
+MASK_TOKEN = '[MASK]'
+MASKING_CHANCE = 0.1
+
+class DataEncodings(Dataset):
+    '''
+    tokenizes, pads, and adds special tokens
+    '''
+    def __init__(
+        self,
+        data: pd.DataFrame,
+        tokenizer,
+        source_max_token_len: int,
+        target_max_token_len: int
+        ):
+        self.tokenizer = t5tokenizer
+        self.data = data
+        self.source_max_token_len = source_max_token_len
+        self.target_max_token_len = target_max_token_len
+
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, index:int):
+        data_row = self.data.iloc[index]
+        # adds a random mask for answer-agnostic qg
+        if np.random.rand() > MASKING_CHANCE:
+            answer = data_row['answer']
+        else:
+            answer = MASK_TOKEN
+            
+        source_encoding = t5tokenizer(
+            f"{answer} {SEP_TOKEN} {data_row['context']}",
+            max_length= self.source_max_token_len,
+            padding='max_length',
+            truncation= True,
+            return_attention_mask=True,
+            return_tensors='pt'
+            )
+    
+        target_encoding = t5tokenizer(
+            f"{data_row['answer']} {SEP_TOKEN} {data_row['question']}",
+            max_length=self.target_max_token_len,
+            padding='max_length',
+            truncation = True,
+            return_attention_mask=True,
+            return_tensors='pt'
+            )
+
+        labels = target_encoding['input_ids']  
+        labels[labels == 0] = -100 # masked
+
+        encodings = dict(
+            answer = data_row['answer'],
+            context = data_row['context'],
+            question = data_row['question'],
+            input_ids = source_encoding['input_ids'].flatten(),
+            attention_mask = source_encoding['attention_mask'].flatten(),
+            labels=labels.flatten()
+        )
+        
+        return encodings
+
+class DataModule(pl.LightningDataModule):
+
+    def __init__(
+        self,
+        train: pd.DataFrame,
+        val: pd.DataFrame,
+        tokenizer,
+        batch_size,
+        source_max_token_len: int,
+        target_max_token_len: int
+        ): 
+        super().__init__()
+        self.batch_size = batch_size
+        self.train = train
+        self.val = val
+        self.tokenizer = t5tokenizer
+        self.source_max_token_len = source_max_token_len
+        self.target_max_token_len = target_max_token_len
+
+    def setup(self):
+        self.train_dataset = DataEncodings(self.train, self.tokenizer, self.source_max_token_len, self.target_max_token_len)
+        self.val_dataset = DataEncodings(self.val, self.tokenizer, self.source_max_token_len, self.target_max_token_len)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
+
+    def val_dataloader(self): 
+        return DataLoader(self.val_dataset, batch_size=batch_size, num_workers=0)
+
+# hyperparameters
+num_epochs = 16
+batch_size = 32
+learning_rate = 0.001
+
+# model 
+class T5Model(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.model = t5model
+        self.model.resize_token_embeddings(len(t5tokenizer)) # resizing after adding new tokens to the tokenizer
+
+    # feed forward pass
+    def forward(self, input_ids, attention_mask, labels=None):
+        output = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+        return output.loss, output.logits
+
+    # train model and compute loss
+    def training_step(self, batch, batch_idx):
+        input_ids = batch['input_ids']
+        attention_mask = batch['attention_mask']
+        labels = batch['labels']
+        loss, output = self(input_ids, attention_mask, labels)
+        self.log('train_loss', loss, prog_bar=True, logger=True, batch_size=batch_size)
+        return loss
+
+    # gets model predictions, returns loss
+    def validation_step(self, batch, batch_idx):
+        input_ids = batch['input_ids']
+        attention_mask = batch['attention_mask']
+        labels = batch['labels']
+        loss, output = self(input_ids, attention_mask, labels)
+        self.log('val_loss', loss, prog_bar=True, logger=True, batch_size=batch_size)
+        return {'val loss': loss}
+    
+    # def validation_epoch_end(self, outputs):
+    #     # outputs = list of dictionaries to print loss
+    #     avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
+    #     tensorboard_logs = {'avg_val_loss': avg_loss}
+    #     return {'val_loss': avg_loss, 'log': tensorboard_logs}
+
+    def configure_optimizers(self):
+        return Adafactor(model.parameters(), scale_parameter=False, relative_step=False, lr=learning_rate) 
+
+def generate(model: T5Model, answer:str, context:str, beams, length, temper) -> str:
+    source_encoding = t5tokenizer(
+        f"{answer} {SEP_TOKEN} {context}",
+        max_length=512,
+        padding='max_length',
+        truncation=True,
+        return_attention_mask=True,
+        add_special_tokens=True,
+        return_tensors='pt'
+    )
+
+    generated_ids=model.model.generate(
+        input_ids=source_encoding['input_ids'],
+        attention_mask=source_encoding['attention_mask'],
+        num_beams=beams,
+        max_length=length,
+        repetition_penalty=2.5,
+        length_penalty=0.8,
+        temperature=temper,
+        early_stopping=True,
+        use_cache=True
+    )
+
+    preds = {
+        t5tokenizer.decode(generated_id, skip_special_tokens=False, clean_up_tokenization_spaces=True)
+        for generated_id in generated_ids
+    }
+
+    return ''.join(preds)
+
+def show_result(generated:str, answer:str, context:str, original_question:str=''):
+
+    regex = r"(?<=>)(.*?)(?=<)"
+    matches = re.findall(regex, generated)
+    matches[1] = matches[1][5:]
+    final = {cat: match.strip() for cat, match in zip(['Answer', 'Question'], matches)}
+    st.title('Context')
+    st.write(context)
+    st.title('Answer')
+    st.write(answer)
+    st.title('Generated')
+    st.write(final)
+    # if original_question:
+    #     printBold('Original Question')
+    #     print(original_question)
+    #     gen = nlp(matches[1])
+    #     ori = nlp(original_question)
+    #     bleu_score = sentence_bleu(matches[1], original_question, smoothing_function=SmoothingFunction().method5)
+    #     cs_score = ori.similarity(gen)
+    #     printBold('Scores')
+    #     print(f"BLEU: {bleu_score}")
+    #     print(f'Cosine Similarity: {cs_score}')
+    #     return bleu_score, cs_score
+
+# streamlit app
+st.title('Question Generation From Text')
+
+with st.form('my_form'):
+    context = st.text_input('Enter a context passage for question generation:', 'The capital of France is Paris.')
+    answer = st.text_input('Give a correct answer, or [MASK] for unsupervised generation:', 'Paris')
+    # question = st.text_input('Question', 'What is the capital of France?')
+    # original_question = st.text_input('Original Question', 'What is the capital of France?')
+    beams = st.sidebar.slider('Beams', min_value=1, max_value=20)
+    length = st.sidebar.slider('Maximum length of generated question', min_value=50, max_value=200)
+    temper = st.sidebar.slider("Temperature", value = 1.0, min_value = 0.0, max_value=1.0, step=0.05)
+    submitted = st.form_submit_button('Generate')
+
+with st.spinner('Loading Model...'):
+    model = T5Model
+    best_model_dir = '../checkpoints/t5-chkpt-v2.ckpt'
+    best_model = model.load_from_checkpoint(best_model_dir)
+    # best_model = model.load_from_checkpoint(callback.best_model_path)
+    best_model.freeze()
+
+with st.spinner('Generating...'):
+    if submitted:
+        
+        generated = generate(best_model, answer, context, beams, length, temper)
+        show_result(generated, answer, context)
+