Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| import pandas as pd | |
| import matplotlib.pyplot as plt # for making figures | |
| # %matplotlib inline | |
| # %config InlineBackend.figure_format = 'retina' | |
| from pprint import pprint | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| data = open('shakespeare2.txt', 'r').read() | |
| unique_chars = list(set(''.join(data))) | |
| unique_chars.sort() | |
| to_string = {i:ch for i, ch in enumerate(unique_chars)} | |
| to_int = {ch:i for i, ch in enumerate(unique_chars)} | |
| class NextChar(nn.Module): | |
| def __init__(self, block_size, vocab_size, emb_dim, hidden_dims): | |
| super().__init__() | |
| self.emb = nn.Embedding(vocab_size, emb_dim) | |
| self.lin1 = nn.Linear(block_size * emb_dim, hidden_dims[0]) | |
| self.lin2 = nn.Linear(hidden_dims[0], hidden_dims[1]) | |
| self.lin3 = nn.Linear(hidden_dims[1], vocab_size) | |
| def forward(self, x): | |
| x = self.emb(x) | |
| x = x.view(x.shape[0], -1) | |
| x = torch.sin(self.lin1(x)) | |
| x = torch.sin(self.lin2(x)) | |
| x = self.lin3(x) | |
| return x | |
| # For context size 5 and embedding size 5 | |
| model_c5_e5 = NextChar(5, len(to_int), 5, [64, 64]) | |
| model_c5_e5.load_state_dict(torch.load("context_5_embedding_5.pth")) | |
| # For context size 5 and embedding size 10 | |
| model_c5_e10 = NextChar(5, len(to_int), 10, [64, 64]) | |
| model_c5_e10.load_state_dict(torch.load("context_5_embedding_10.pth")) | |
| # For context size 7 and embedding size 5 | |
| model_c7_e5 = NextChar(7, len(to_int), 5, [64, 64]) | |
| model_c7_e5.load_state_dict(torch.load("context_7_embedding_5.pth")) | |
| # For context size 7 and embedding size 10 | |
| model_c7_e10 = NextChar(7, len(to_int), 10, [64, 64]) | |
| model_c7_e10.load_state_dict(torch.load("context_7_embedding_10.pth")) | |
| random_seed = 3 | |
| g = torch.Generator() | |
| g.manual_seed(random_seed) | |
| torch.manual_seed(random_seed) | |
| def generate_name(model,sentence, itos, stoi, block_size, max_len=10): | |
| original_sentence = sentence | |
| if len(sentence) < block_size: | |
| sentence = " " * (block_size - len(sentence)) + sentence | |
| using_for_predicction = sentence[-block_size:].lower() | |
| context = [stoi[word] for word in using_for_predicction] | |
| prediction = "" | |
| for i in range(max_len): | |
| x = torch.tensor(context).view(1, -1).to(device) | |
| print(type(model)) | |
| y_pred = model(x) | |
| ix = torch.distributions.categorical.Categorical(logits=y_pred).sample().item() | |
| ch = itos[ix] | |
| prediction += ch | |
| context = context[1:] + [ix] | |
| return original_sentence + prediction | |
| # Streamlit app | |
| st.title("Next K Text Generation with MLP") | |
| st.sidebar.title("Settings") | |
| input_string = st.sidebar.text_input("Input String") | |
| nextk = st.sidebar.number_input("Next K Tokens", min_value=1, max_value=500, value=150) | |
| block_size = st.select_slider("Block Size", options=[5,7], value=5) | |
| embedding_size = st.select_slider("Embedding Size", options=[5,10], value=5) | |
| if st.sidebar.button("Generate Text"): | |
| if block_size == 5: | |
| context = input_string | |
| if embedding_size == 5: | |
| generated_text = generate_name(model_c5_e5,context, to_string, to_int, 5, max_len=nextk) | |
| else: | |
| generated_text = generate_name(model_c5_e10,context, to_string, to_int, 5, max_len=nextk) | |
| elif block_size == 7: | |
| context = input_string | |
| if embedding_size == 10: | |
| generated_text = generate_name(model_c7_e10, context, to_string, to_int, 7 ,max_len=nextk) | |
| else: | |
| generated_text = generate_name(model_c7_e5, context, to_string, to_int, 7, max_len=nextk) | |
| st.write("Generated Text:") | |
| st.write(generated_text) | |