Spaces:
Sleeping
Sleeping
File size: 3,261 Bytes
47d6804 34a1570 47d6804 34a1570 47d6804 fce3e3b 47d6804 |
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 104 105 106 107 108 109 110 111 112 113 114 115 |
import streamlit as st
import torch
from pathlib import Path
import math
from dataclasses import dataclass
import torch.nn as nn
import torch.nn.functional as F
from src.gpt_base import GPT
import json
from huggingface_hub import hf_hub_download
from tqdm import tqdm
# Config class for model parameters
@dataclass
class GPTConfig:
block_size: int = 1024 # max sequence length
vocab_size: int = 65
num_layer: int = 12 # number of layers
num_head: int = 12 # number of heads
emb_dim: int = 768 # embedding dimension
dropout: float = 0.1 # dropout rate
# Copy all the model classes (GPT, MultiHeadAttention, FeedForward, TransformerBlock) here
# [Previous model code goes here]
# Load stoi and itos from docs
with open("docs/stoi.json") as f:
stoi = json.load(f)
with open("docs/itos.json") as f:
itos = json.load(f)
# Encoding/Decoding functions
def encode(s):
return [stoi[c] for c in s]
def decode(l):
return "".join([itos[i] for i in l])
def predict_next_word(text, model, seq_len=50):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
for _ in tqdm(range(seq_len)):
xb = torch.tensor(encode(text)).unsqueeze(0).to(device)
yb = model(xb)
next_word = yb[0, -1].argmax().item()
text += itos[str(next_word)]
return text
# Streamlit app
st.title("GPT Text Generation")
# Add some usage instructions
st.markdown(
"""
### How to use:
1. Enter your text prompt in the text box below
2. Adjust the sequence length using the slider
3. Click 'Generate Text' to see the model's output
Note: Longer sequence lengths will take more time to generate.
"""
)
# Input text box
input_text = st.text_area("Enter your text prompt:", height=100)
# Sequence length slider
seq_length = st.slider(
"Select sequence length for prediction:",
min_value=50,
max_value=500,
value=200,
step=50,
)
# Model loading and prediction
if st.button("Generate Text"):
if input_text:
try:
# Initialize model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
config = GPTConfig()
model = GPT(config)
model = model.to(device)
# Load checkpoint
# checkpoint_path = "/Users/aditya/Documents/self_learning/ERA V3/week 12/model artifacts/gpt_model_and_loss.pth"
model_repo = "Adityak204/JuliusCaesarGPT"
model_filename = "gpt_model_and_loss.pth"
checkpoint_path = hf_hub_download(
repo_id=model_repo, filename=model_filename
)
with st.spinner("Loading model and generating text..."):
_dict = torch.load(checkpoint_path, map_location=device)
model_state_dict = _dict["model_state_dict"]
model.load_state_dict(model_state_dict)
# Generate text
generated_text = predict_next_word(input_text, model, seq_length)
# Display results
st.subheader("Generated Text:")
st.write(generated_text)
except Exception as e:
st.error(f"An error occurred: {str(e)}")
else:
st.warning("Please enter some text first!")
|