import shutil
import requests
import pandas as pd
import numpy as np
import seaborn as sns 
import matplotlib.pyplot as plt
import torch.nn.functional as F
import torch
import streamlit as st
from torch import nn
from torch.optim import Adam
from tqdm import tqdm
from bs4 import BeautifulSoup 
from sklearn.model_selection import train_test_split
from collections import defaultdict
from transformers import RobertaTokenizerFast, RobertaForQuestionAnswering

#加載模型參數
def load_ckp(checkpoint_fpath, model, optimizer):
    """
    checkpoint_path: path to save checkpoint
    model: model that we want to load checkpoint parameters into       
    optimizer: optimizer we defined in previous training
    """
    # load check point
    # 如果使用cpu則後面需加上 map_location=torch.device('cpu')
    checkpoint = torch.load(checkpoint_fpath,map_location=torch.device('cpu'))
    # initialize state_dict from checkpoint to model
    model.load_state_dict(checkpoint['state_dict'])
    # initialize optimizer from checkpoint to optimizer
    optimizer.load_state_dict(checkpoint['optimizer'])
    # initialize valid_loss_min from checkpoint to valid_loss_min
    valid_loss_min = checkpoint['valid_loss_min']
    # return model, optimizer, epoch value, min validation loss 
    return model, optimizer, checkpoint['epoch'], valid_loss_min

#使用GPU
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

#預訓練模組
tokenizer = RobertaTokenizerFast.from_pretrained("roberta-base")
model = RobertaForQuestionAnswering.from_pretrained("roberta-base")

#優化器
LEARNING_RATE=3e-5
optimizer = torch.optim.Adam(params =  model.parameters(), lr=LEARNING_RATE)


#加載最優參數
load_ckp('./model/best_model.pt', model, optimizer)


#預測函數
def Predict(question,context):
  inputs = tokenizer(question,context,max_length= 512,padding='max_length',return_offsets_mapping=True,return_tensors="pt")
  input_ids=inputs['input_ids'].to(device)
  attention_mask =inputs['attention_mask'].to(device)
  outputs = model(input_ids.reshape(1,512),attention_mask.reshape(1,512))
  answer_start_index = outputs.start_logits.argmax(dim=1)
  answer_end_index = outputs.end_logits.argmax(dim=1)

  #predict answer
  predict_answer = tokenizer.decode(inputs['input_ids'].flatten()[answer_start_index.item() : answer_end_index.item()])
  return predict_answer.strip()


# qa = load_ckp('./model/best_model.pt', model, optimizer)
def main():
    st.title("Question Answering")

    with st.form("text_field"):
        question = st.text_input("Enter some question :")
        context = st.text_area("Enter some context :")
        
        # clicked==True only when the button is clicked
        clicked = st.form_submit_button("Submit")
        if clicked:        
          results = Predict(question = question, context = context)
          st.json(results)


if __name__ == "__main__":
    main()