parkermoe/RecipeGenerator

1import torch2import torch.nn as nn3import os4import pickle5from torch.functional import F6import numpy as np7import gradio as gr8import torchtext9 10#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")11device = torch.device('cpu')12VOCAB_SIZE = 1000013MAX_LEN = 20014EMBEDDING_DIM = 10015N_UNITS = 12816VALIDATION_SPLIT = 0.217SEED = 4218LOAD_MODEL = False19BATCH_SIZE = 12820EPOCHS = 2521# loading model from checkpoint22class LSTMModel(nn.Module):23    def __init__(self, vocab_size, embedding_dim, hidden_dim):24        super(LSTMModel, self).__init__()25        self.embedding = nn.Embedding(vocab_size, embedding_dim)26        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)27        self.fc = nn.Linear(hidden_dim, vocab_size)28        self.log_softmax = nn.LogSoftmax(dim=2)29 30    def forward(self, x):31        x = self.embedding(x)32        x, _ = self.lstm(x)33        x = self.fc(x)34        return self.log_softmax(x)35 36# loading model from checkpoint37model = LSTMModel(VOCAB_SIZE, EMBEDDING_DIM, N_UNITS).to(device)38 39device = 'cpu'40 41checkpoint_path = 'recipe_generator_LSTM.pth'42checkpoint = torch.load(checkpoint_path, map_location=device)43model.load_state_dict(checkpoint)44 45print('Loaded model from checkpoint')46 47def load_vocab(file_path):48    file_path = os.path.join(file_path)49    with open(file_path, 'rb') as input:50        vocab = pickle.load(input)51    print(f"Vocabulary loaded from {file_path}")52    return vocab53 54vocab = load_vocab('vocab.pkl')55 56 57 58 59class TextGenerator:60    def __init__(self, vocab, top_k=10):61        self.vocab = vocab62        self.top_k = top_k63 64    def sample_from(self, logits, temperature):65        probs = F.softmax(logits / temperature, dim=-1).cpu().numpy()66        return np.random.choice(len(probs), p=probs)67 68    def generate(self, model, device, start_prompt, max_tokens, temperature):69        model.eval()70 71        tokens = [self.vocab.get_stoi()[token] for token in start_prompt.split()]72        tokens = torch.LongTensor(tokens).unsqueeze(0).to(device)73 74        with torch.no_grad():75            for _ in range(max_tokens):76                output = model(tokens)77                next_token_logits = output[0, -1, :]78                next_token = self.sample_from(next_token_logits, temperature)79                tokens = torch.cat([tokens, torch.LongTensor([[next_token]]).to(device)], dim=1)80 81        generated_tokens = [token for token in tokens[0] if self.vocab.get_itos()[token] != '<pad>']82        generated_text = ' '.join(self.vocab.get_itos()[token] for token in generated_tokens)83        return generated_text84    85text_generator = TextGenerator(vocab=vocab, top_k=10)86generated_text = text_generator.generate(model=model, device=device, start_prompt="recipe for", max_tokens=100, temperature=0.5)87 88print(f"\nGenerated Text: {generated_text}")89 90 91 92def generate_recipe():93    return text_generator.generate(model=model, device=device, start_prompt="recipe for", max_tokens=100, temperature=0.5)94 95iface = gr.Interface(96    fn=generate_recipe, 97    inputs=[], 98    outputs="text",99    title="Recipe Generator",100    description="This is a LSTM based Recurrent Neural Network trained to generate recipes. Press submit to generate a new recipe that can sometimes provide humor!",101)102 103iface.launch()