Spaces:
Sleeping
Sleeping
File size: 2,477 Bytes
dd7ff72 039c272 dd7ff72 039c272 dd7ff72 039c272 dd7ff72 039c272 dd7ff72 039c272 dd7ff72 039c272 dd7ff72 039c272 dd7ff72 039c272 dd7ff72 039c272 dd7ff72 039c272 dd7ff72 039c272 dd7ff72 039c272 dd7ff72 039c272 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import os
os.environ["HF_HOME"] = os.path.abspath("./.hf_cache")
import torch
import pickle
import yaml
from datasets import load_dataset
from model import load_model_processor
config = yaml.safe_load(open("./config.yaml", "r"))["data_location"]
class VQADataset(torch.utils.data.Dataset):
def __init__(self, dataset, processor):
self.dataset = dataset
self.processor = processor
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
question = self.dataset[idx]["question"]
answer = self.dataset[idx]["answer"]
image = self.dataset[idx]["image"]
image = image.convert("RGB")
text = question
encoding = self.processor(
image, text, padding="max_length", truncation=True, return_tensors="pt"
)
labels = self.processor.tokenizer.encode(
answer,
max_length=128,
padding="max_length",
truncation=True,
pad_to_max_length=True,
return_tensors="pt",
)
encoding["labels"] = labels
for k, v in encoding.items():
encoding[k] = v.squeeze()
return encoding
if __name__ == "__main__":
_, processor = load_model_processor()
from datasets import load_from_disk
print("Loading VQA dataset................")
# Load dataset that was pre-saved to disk natively
data = load_from_disk("./data/bronze")
train_data = data["train"]
test_data = data["test"]
print(
"VQA dataset loaded successfully!!! lenght of train data is ",
len(train_data),
" and test data is ",
len(test_data),
)
save_dir = "./data/silver"
os.makedirs(save_dir, exist_ok=True)
print("Processesing data to save in ../data/silver")
train_dataset = VQADataset(dataset=train_data, processor=processor)
test_dataset = VQADataset(dataset=test_data, processor=processor)
print(f"Data processed successfully !!! ")
print(
f"Saving to {os.path.join(save_dir, 'train_dataset.pkl')} and {os.path.join(save_dir, 'test_dataset.pkl')}"
)
with open(os.path.join(save_dir, "train_dataset.pkl"), "wb") as f:
pickle.dump(train_dataset, f)
with open(os.path.join(save_dir, "test_dataset.pkl"), "wb") as f:
pickle.dump(test_dataset, f)
print(
f"Processed data , saved to {config['train_processed_data']} and {config['test_processed_data']}"
)
|