add cache dir

app.py

@@ -8,7 +8,8 @@ from model_loader import *
 
 
 # load dataset
-ds = load_dataset("test")
+#ds = load_dataset("test")
+ds = load_dataset("HuggingFaceM4/VQAv2", split="validation", cache_dir="cache", streaming=False)
 
 # define selector
 model_name = st.sidebar.selectbox(

model_loader.py

@@ -8,7 +8,6 @@ import requests
 from transformers import ViltProcessor, ViltForQuestionAnswering
 from transformers import AutoProcessor, AutoModelForCausalLM
 from transformers import BlipProcessor, BlipForQuestionAnswering
-from nltk.corpus import wordnet
 
 import os
 import requests
@@ -25,7 +24,6 @@ import torchvision.transforms as transforms
 from transformers import VisualBertForMultipleChoice, VisualBertForQuestionAnswering, BertTokenizerFast, AutoTokenizer, ViltForQuestionAnswering
 
 from PIL import Image
-from nltk.corpus import wordnet
 import time
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -50,14 +48,15 @@ def load_model(name):
 
     return (processor, model)
 
-
+'''
 def load_dataset(type):
     if type == "train":
-        return load_dataset("HuggingFaceM4/VQAv2", split="train", streaming=False)
+        return load_dataset("HuggingFaceM4/VQAv2", split="train", cache_dir="cache", streaming=False)
     elif type == "test":
-        return load_dataset("HuggingFaceM4/VQAv2", split="validation", streaming=False)
+        return load_dataset("HuggingFaceM4/VQAv2", split="validation", cache_dir="cache", streaming=False)
     else:
         raise ValueError("invalid dataset: ", type)
+'''
 
 
 def tokenize_function(examples, processor):