Update app.py

app.py

@@ -2,10 +2,12 @@ import gradio as gr
 from transformers import BertForQuestionAnswering
 from transformers import BertTokenizerFast
 import torch
+from datasets import load_dataset
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
 model = BertForQuestionAnswering.from_pretrained("CountingMstar/ai-tutor-bert-model").to(device)
+dataset = load_dataset("CountingMstar/AI_term", split="train")
 
 def get_prediction(context, question):
     inputs = tokenizer.encode_plus(question, context, return_tensors='pt').to(device)