# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="crPVEDSN71sl"
# <a href="https://colab.research.google.com/github/khamies//LSTM-Language-Generator/blob/master/play_with_model.ipynb" 
# target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
#

# + colab={"base_uri": "https://localhost:8080/"} id="akHUDpl5RRWk" cellView="form" outputId="ee99fdd9-0c33-43b4-9f57-af543e6a724a"
#@title Download data and files.
# !git clone https://github.com/Khamies/LSTM-Language-Generator.git
import os 
os.chdir("LSTM-Language-Generator")

# + id="bsX47UMWT0NT"
import torch
from data.ptb import PTB
from model import LSTM_Language

# + id="u6QyBIaJWGLI"
# Settings

torch.manual_seed(1000)
device = "cuda" if torch.cuda.is_available() else "cpu"

batch_size = 32
bptt = 60
lr = 0.001

embed_size = 300
hidden_size = 256
latent_size = 16
lstm_layer=1

# + [markdown] id="GcXo1SRDfIVm"
# ## Load the data

# + id="k2uVseOtT_4J"
# Load the data
train_data = PTB(data_dir="./data", split="train", create_data= False, max_sequence_length= bptt)
test_data = PTB(data_dir="./data", split="test", create_data= False, max_sequence_length= bptt)
valid_data = PTB(data_dir="./data", split="valid", create_data= False, max_sequence_length= bptt)

# Batchify the data
train_loader = torch.utils.data.DataLoader( dataset= train_data, batch_size= batch_size, shuffle= True)
test_loader = torch.utils.data.DataLoader( dataset= test_data, batch_size= batch_size, shuffle= True)
valid_loader = torch.utils.data.DataLoader( dataset= valid_data, batch_size= batch_size, shuffle= True)

# + [markdown] id="mZ2jOq3se8-x"
# ## Load the model

# + colab={"base_uri": "https://localhost:8080/"} id="jw2_Fc_9Vhh5" outputId="1e0886bb-8969-4d43-efbd-082656d391cb"
vocab_size = train_data.vocab_size
model = LSTM_Language(vocab_size = vocab_size, embed_size = embed_size, hidden_size = hidden_size, latent_size = latent_size).to(device)

checkpoint = torch.load("models/LSTM_lang.pt")
model.load_state_dict(checkpoint["model"])

# + [markdown] id="7s4CNs1J4otu"
# ##Sample Generation

# + colab={"base_uri": "https://localhost:8080/"} id="PKhUHL_7Vg5X" cellView="code" outputId="f3b2e51f-3e87-4eb6-f26d-6cd1aa70a639"
sos = "<sos>"
sample = model.inference(10 , sos)


print(sample)


# + id="a_XvIj4l4QMi"


StarcoderdataJupyterScriptsDedupFiltered