app.py

# coding=utf8

from transformers import AutoModel, AutoTokenizer, AutoConfig
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import streamlit as st
import gdown
import numpy as np
import pandas as pd
import collections
from string import punctuation


class CONFIG:
    #model params
    model = 'deepset/xlm-roberta-large-squad2'
    max_input_length = 384 #Hyperparameter to be tuned, following the guide from huggingface
    doc_stride = 128  #Hyperparameter to be tuned, following the guide from huggingface
    model_checkpoint = "pytorch_model.pth"
    trained_model_url = 'https://drive.google.com/uc?id=16Vp918RglyLEFEyDlFuRD1HeNZ8SI7P5'
    trained_model_output_fp = 'trained_pytorch.pth'
    sample_df_fp = "sample_qa.json"

# model class
class ChaiModel(nn.Module):
    def __init__(self, model_config):
        super(ChaiModel, self).__init__()
        self.backbone = AutoModel.from_pretrained(CONFIG.model)
        self.linear = nn.Linear(model_config.hidden_size, 2)
        
    def forward(self, input_ids, attention_mask):
        model_output = self.backbone(input_ids, attention_mask=attention_mask)
        sequence_output = model_output[0] # (batchsize, sequencelength, hidden_dim)
        
        qa_logits = self.linear(sequence_output) # (batchsize, sequencelength, 2)
        start_logit, end_logit = qa_logits.split(1, dim=-1) #  (batchsize, sequencelength), 1), (batchsize, sequencelength, 1)
        start_logits = start_logit.squeeze(-1) # remove last dim (batchsize, sequencelength)
        end_logits = end_logit.squeeze(-1)    #remove last dim (batchsize, sequencelength)
        
        return start_logits, end_logits # (2,batchsize, sequencelength)

# dataset class
class ChaiDataset(Dataset):
    def __init__(self, dataset, is_train=True):
        super(ChaiDataset, self).__init__()
        self.dataset = dataset #list of features
        self.is_train= is_train
        
    def __len__(self):
        return len(self.dataset)
    
    def __getitem__(self, index):
        features = self.dataset[index]
        if self.is_train:
            return {
                'input_ids': torch.tensor(features['input_ids'], dtype=torch.long),
                'attention_mask': torch.tensor(features['attention_mask'], dtype=torch.long),
                'offset_mapping':torch.tensor(features['offset_mapping'], dtype=torch.long),
                'start_position':torch.tensor(features['start_position'], dtype=torch.long),
                'end_position':torch.tensor(features['end_position'], dtype=torch.long)
            }
        else:
            return {
                'input_ids': torch.tensor(features['input_ids'], dtype=torch.long),
                'attention_mask': torch.tensor(features['attention_mask'], dtype=torch.long),
                'offset_mapping':torch.tensor(features['offset_mapping'], dtype=torch.long),
                'sequence_ids':features['sequence_ids'],
                'id':features['example_id'],
                'context':features['context'],
                'question':features['question']
            }
    
def break_long_context(df, tokenizer, train=True):
    if train: 
        n_examples = len(df)
        full_set = []
        for i in range(n_examples):
            row = df.iloc[i]
            # tokenizer parameters can be found here 
            # https://huggingface.co/transformers/internal/tokenization_utils.html#transformers.tokenization_utils_base.PreTrainedTokenizerBase
            tokenized_examples = tokenizer(row['question'],
                                          row['context'],
                                          padding='max_length',
                                          max_length=CONFIG.max_input_length, 
                                          truncation='only_second',
                                          stride=CONFIG.doc_stride,
                                          return_overflowing_tokens=True, #returns the number of over flow
                                          return_offsets_mapping=True     #returns the BPE mapping to the original word
                                          ) 
            
            # tokenized_example keys
            #'input_ids', 'attention_mask', 'offset_mapping', 'overflow_to_sample_mapping'
            sample_mappings = tokenized_examples.pop("overflow_to_sample_mapping")
            offset_mappings = tokenized_examples.pop("offset_mapping")
            
            final_examples = [] 
            n_sub_examples = len(sample_mappings)
            for j in range(n_sub_examples):
                input_ids = tokenized_examples["input_ids"][j]
                attention_mask = tokenized_examples["attention_mask"][j]
                
                sliced_text = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(input_ids))
                final_example = dict(input_ids = input_ids, 
                                     attention_mask = attention_mask,
                                     sliced_text = sliced_text,
                                     offset_mapping=offset_mappings[j],
                                     fold=row['fold'])
                
                
                
                # Most of the time cls_index is 0
                cls_index = input_ids.index(tokenizer.cls_token_id)
                # None, 0, 0, .... None, None, 1, 1,.....
                sequence_ids = tokenized_examples.sequence_ids(j)
                
                sample_index = sample_mappings[j]
                offset_map = offset_mappings[j]
                
                if np.isnan(row["answer_start"]) : # if no answer, start and end position is cls_index
                    final_example['start_position'] = cls_index
                    final_example['end_position'] = cls_index
                    final_example['tokenized_answer'] = ""
                    final_example['answer_text'] = ""
                else:
                    start_char  = row["answer_start"]
                    end_char  = start_char + len(row["answer_text"])
                    
                    token_start_index = sequence_ids.index(1)
                    token_end_index = len(sequence_ids)- 1 - (sequence_ids[::-1].index(1))
                    
                    if not (offset_map[token_start_index][0]<=start_char and offset_map[token_end_index][1] >= end_char):
                        final_example['start_position'] = cls_index
                        final_example['end_position'] = cls_index
                        final_example['tokenized_answer'] = ""
                        final_example['answer_text'] = ""
                    else:
                        #Move token_start_index to the correct context index
                        while token_start_index < len(offset_map) and offset_map[token_start_index][0] <= start_char:
                            token_start_index +=1
                        final_example['start_position'] = token_start_index -1
                        
                        while offset_map[token_end_index][1] >= end_char: #Take note that we will want the end_index inclusively, we will need to slice properly later
                            token_end_index -=1
                        final_example['end_position'] = token_end_index + 1   
                        tokenized_answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(input_ids[final_example['start_position']:final_example['end_position']+1]))
                        final_example['tokenized_answer'] = tokenized_answer
                        final_example['answer_text'] = row['answer_text']
                        
                final_examples.append(final_example)
            full_set += final_examples
            
    else:
        n_examples = len(df)
        full_set = []
        for i in range(n_examples):
            row = df.iloc[i]
            tokenized_examples = tokenizer(row['question'], 
                                          row['context'],
                                          padding='max_length',
                                          max_length=CONFIG.max_input_length, 
                                          truncation='only_second',
                                          stride=CONFIG.doc_stride,
                                          return_overflowing_tokens=True, #returns the number of over flow
                                          return_offsets_mapping=True     #returns the BPE mapping to the original word
                                          )
            
            sample_mappings = tokenized_examples.pop("overflow_to_sample_mapping")
            offset_mappings = tokenized_examples.pop("offset_mapping")
            n_sub_examples = len(sample_mappings)
            
            final_examples = []
            for j in range(n_sub_examples):
                input_ids = tokenized_examples["input_ids"][j]
                attention_mask = tokenized_examples["attention_mask"][j]
                
                final_example = dict(
                    input_ids = input_ids, 
                    attention_mask = attention_mask,
                    offset_mapping=offset_mappings[j],
                    example_id = row['id'],
                    context = row['context'],
                    question = row['question'],
                    sequence_ids = [0 if value is None else value for value in tokenized_examples.sequence_ids(j)]  
                )
                
                final_examples.append(final_example)
            full_set += final_examples
        return full_set

def postprocess_qa_predictions(examples, features, raw_predictions, n_best_size = 20, max_answer_length = 30):
    all_start_logits, all_end_logits = raw_predictions
    
    example_id_to_index = {k: i for i, k in enumerate(examples["id"])}
    features_per_example = collections.defaultdict(list)
    for i, feature in enumerate(features):
        features_per_example[example_id_to_index[feature["example_id"]]].append(i)

    predictions = collections.OrderedDict()

    print(f"Post-processing {len(examples)} example predictions split into {len(features)} features.")

    for example_index, example in examples.iterrows():
        feature_indices = features_per_example[example_index]

        min_null_score = None
        valid_answers = []
        
        context = example["context"]
        for feature_index in feature_indices:
            start_logits = all_start_logits[feature_index]
            end_logits = all_end_logits[feature_index]

            sequence_ids = features[feature_index]["sequence_ids"]
            context_index = 1

            features[feature_index]["offset_mapping"] = [
                (o if sequence_ids[k] == context_index else None)
                for k, o in enumerate(features[feature_index]["offset_mapping"])
            ]
            offset_mapping = features[feature_index]["offset_mapping"]
            cls_index = features[feature_index]["input_ids"].index(tokenizer.cls_token_id)
            feature_null_score = start_logits[cls_index] + end_logits[cls_index]
            if min_null_score is None or min_null_score < feature_null_score:
                min_null_score = feature_null_score

            start_indexes = np.argsort(start_logits)[-1 : -n_best_size - 1 : -1].tolist()
            end_indexes = np.argsort(end_logits)[-1 : -n_best_size - 1 : -1].tolist()
            for start_index in start_indexes:
                for end_index in end_indexes:
                    if (
                        start_index >= len(offset_mapping)
                        or end_index >= len(offset_mapping)
                        or offset_mapping[start_index] is None
                        or offset_mapping[end_index] is None
                    ):
                        continue
                    # Don't consider answers with a length that is either < 0 or > max_answer_length.
                    if end_index < start_index or end_index - start_index + 1 > max_answer_length:
                        continue

                    start_char = offset_mapping[start_index][0]
                    end_char = offset_mapping[end_index][1]
                    valid_answers.append(
                        {
                            "score": start_logits[start_index] + end_logits[end_index],
                            "text": context[start_char: end_char]
                        }
                    )
        
        if len(valid_answers) > 0:
            best_answer = sorted(valid_answers, key=lambda x: x["score"], reverse=True)[0]
        else:
            best_answer = {"text": "", "score": 0.0}
        
        predictions[example["id"]] = best_answer["text"]
        
        
    return predictions

def download_finetuned_model():
    gdown.download(url=CONFIG.trained_model_url, output=CONFIG.trained_model_output_fp, quiet=False)
    
def get_prediction(context:str, question:str, model, tokenizer) -> str:
    # convert to dataframe format to make it consistent with training way
    test_df = pd.DataFrame({"id":[1], "context":[context.strip()], "question":[question.strip()]})
    test_set = break_long_context(test_df, tokenizer, train=False)
    
    #create dataset and dataloader of batch 1 to prevent OOM
    test_dataset = ChaiDataset(test_set, is_train=False)
    test_dataloader = DataLoader(test_dataset,
                                 batch_size=1,
                                 shuffle=False,
                                 drop_last=False
                                 )

    #main prediction function
    start_logits =[]
    end_logits=[]

    for features in test_dataloader:
        input_ids = features['input_ids']
        attention_mask = features['attention_mask']
        with torch.no_grad():
            start_logit, end_logit = model(input_ids, attention_mask) #(batch, 384,1) , (batch, 384,1)
            start_logits.append(start_logit.to("cpu").numpy())
            end_logits.append(end_logit.to("cpu").numpy())

    start_logits, end_logits = np.vstack(start_logits),  np.vstack(end_logits)
    
    predictions = postprocess_qa_predictions(test_df, test_set, (start_logits, end_logits))
    predictions = list(predictions.items())[0][1]
    predictions = predictions.strip(punctuation)

    return predictions

@st.cache(allow_output_mutation=True)
def load_model():
    gdown.download(url=CONFIG.trained_model_url, output=CONFIG.trained_model_output_fp, quiet=False)
    print("Downloaded pretrained model")
    config = AutoConfig.from_pretrained(CONFIG.model)
    model = ChaiModel(config)
    model.load_state_dict(torch.load(CONFIG.trained_model_output_fp, map_location=torch.device('cpu')))
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(CONFIG.model)
    sample_df = pd.read_json(CONFIG.sample_df_fp)
    return model, tokenizer, sample_df



model, tokenizer, sample_df =  load_model()


## initialize session_state 
if "context" not in st.session_state:
    st.session_state["context"] = ""
if "question" not in st.session_state:
    st.session_state['question'] = ""
if "answer" not in st.session_state:
    st.session_state['answer'] = ""


## Layout
st.sidebar.title("Hindi/Tamil Extractive Question Answering")
st.sidebar.markdown("---")
random_button = st.sidebar.button("Random")
st.sidebar.write("Randomly Generates a Hindi/Tamil Context and Question")
st.sidebar.markdown("---")
answer_button = st.sidebar.button("Answer!")

if random_button:
    sample = sample_df.sample(1)
    st.session_state['context'] = sample['context'].item()
    st.session_state['question'] = sample['question'].item()
    st.session_state['answer'] = ""

if answer_button:
    # if question or context is empty text
    if len(st.session_state['context']) == 0 or len(st.session_state['question']) ==0:
        st.session_state['answer'] = " "
    else:
        st.session_state['answer'] = get_prediction(st.session_state['context'], st.session_state['question'], model, tokenizer)
    

st.session_state["context"] = st.text_area("Context", value=st.session_state['context'], height=300)

with st.container():
    col_1, col_2 = st.columns(2)
    with col_1:
        st.session_state['question'] = st.text_area("Question", value=st.session_state['question'], height=200)

    with col_2:
        st.text_area("Answer", value=st.session_state['answer'], height=200)