Create predict.py

predict.py

@@ -0,0 +1,84 @@
+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()
+