INNER CODE UNIT · Python

load_vocab

graviraja/100-Days-of-NLP · applications/generation/image_captioning/latex_app/app.py:55

def load_vocab():
    with open(os.path.join("model", 'vocab.pkl'), 'rb') as f:
        vocab = pickle.load(f)
    return vocab

@st.cache
def load_model():
    embed_dim = 256
    encoder_dim = 512
    decoder_dim = 512
    attention_dim = 512
    dropout = 0.3
    vocab = load_vocab()
    encoder = Encoder(encoder_dim).to(device)
    decoder = Decoder(len(vocab), embed_dim, decoder_dim, attention_dim, encoder_dim, dropout).to(device)
    model = Img2LaTeX(encoder, decoder, encoder_dim, decoder_dim, device).to(device)
    model.load_state_dict(torch.load(os.path.join("model", 'model.ckpt'), map_location=torch.device('cpu')))
    return model, vocab

View source record →