Create vocab.py

vocab.py

@@ -0,0 +1,45 @@
+import json
+import os
+import re
+from collections import defaultdict
+import glob
+import numpy as np
+import time
+
+import torch
+import torch.nn as nn
+import torchvision.models as models
+import torch.nn.functional as F
+from torch import optim
+from torch.utils.data import Dataset
+from torchvision import transforms
+from torch.utils.data import DataLoader
+
+from PIL import Image
+
+class Vocabulary:
+    def __init__(self, vocabulary_file_path):
+        #Initialize the Vocabulary object.
+        # Load vocabulary from the provided file path
+        self.vocabulary = self._load_vocabulary(vocabulary_file_path)
+        # Create a mapping from words to indices
+        self.vocabulary2idx = {word: idx for idx, word in enumerate(self.vocabulary)}
+        # Store the total size of the vocabulary
+        self.vocabulary_size = len(self.vocabulary)
+
+    def _load_vocabulary(self, vocabulary_file_path):
+        #Load vocabulary from a file.
+        with open(vocabulary_file_path, 'r') as file:
+            # Read each line, strip extra whitespace, and return as a list
+            vocabulary = [line.strip() for line in file]
+        return vocabulary
+
+    def word2idx(self, word):
+        #Convert a word to its corresponding index.
+        # Return the index of the word or the index of '<unk>' if the word is not in the vocabulary
+        return self.vocabulary2idx.get(word, self.vocabulary2idx.get('<unk>'))
+
+    def idx2word(self, idx):
+        #Convert an index back to its corresponding word.
+        return self.vocabulary[idx]
+