Spaces:
Sleeping
Sleeping
File size: 3,265 Bytes
5bf5b48 ca8df97 5bf5b48 13d544c 5bf5b48 f671f31 13d544c 5bf5b48 29d0dff 5bf5b48 8d4900c 5bf5b48 253aeba 5bf5b48 f46039b 5bf5b48 0123838 |
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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import torch
import torch.nn as nn
import os
import pickle
from torch.functional import F
import numpy as np
import gradio as gr
import torchtext
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device('cpu')
VOCAB_SIZE = 10000
MAX_LEN = 200
EMBEDDING_DIM = 100
N_UNITS = 128
VALIDATION_SPLIT = 0.2
SEED = 42
LOAD_MODEL = False
BATCH_SIZE = 128
EPOCHS = 25
# loading model from checkpoint
class LSTMModel(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim):
super(LSTMModel, self).__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim)
self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
self.fc = nn.Linear(hidden_dim, vocab_size)
self.log_softmax = nn.LogSoftmax(dim=2)
def forward(self, x):
x = self.embedding(x)
x, _ = self.lstm(x)
x = self.fc(x)
return self.log_softmax(x)
# loading model from checkpoint
model = LSTMModel(VOCAB_SIZE, EMBEDDING_DIM, N_UNITS).to(device)
device = 'cpu'
checkpoint_path = 'recipe_generator_LSTM.pth'
checkpoint = torch.load(checkpoint_path, map_location=device)
model.load_state_dict(checkpoint)
print('Loaded model from checkpoint')
def load_vocab(file_path):
file_path = os.path.join(file_path)
with open(file_path, 'rb') as input:
vocab = pickle.load(input)
print(f"Vocabulary loaded from {file_path}")
return vocab
vocab = load_vocab('vocab.pkl')
class TextGenerator:
def __init__(self, vocab, top_k=10):
self.vocab = vocab
self.top_k = top_k
def sample_from(self, logits, temperature):
probs = F.softmax(logits / temperature, dim=-1).cpu().numpy()
return np.random.choice(len(probs), p=probs)
def generate(self, model, device, start_prompt, max_tokens, temperature):
model.eval()
tokens = [self.vocab.get_stoi()[token] for token in start_prompt.split()]
tokens = torch.LongTensor(tokens).unsqueeze(0).to(device)
with torch.no_grad():
for _ in range(max_tokens):
output = model(tokens)
next_token_logits = output[0, -1, :]
next_token = self.sample_from(next_token_logits, temperature)
tokens = torch.cat([tokens, torch.LongTensor([[next_token]]).to(device)], dim=1)
generated_tokens = [token for token in tokens[0] if self.vocab.get_itos()[token] != '<pad>']
generated_text = ' '.join(self.vocab.get_itos()[token] for token in generated_tokens)
return generated_text
text_generator = TextGenerator(vocab=vocab, top_k=10)
generated_text = text_generator.generate(model=model, device=device, start_prompt="recipe for", max_tokens=100, temperature=0.5)
print(f"\nGenerated Text: {generated_text}")
def generate_recipe():
return text_generator.generate(model=model, device=device, start_prompt="recipe for", max_tokens=100, temperature=0.5)
iface = gr.Interface(
fn=generate_recipe,
inputs=[],
outputs="text",
title="Recipe Generator",
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!",
)
iface.launch() |