| import streamlit as st |
| import torch, tiktoken, os |
| from gpt_model import build_model |
|
|
| CKPT = "ckpt_best.pt" |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
| N_LAYER, N_HEAD, N_EMBD, BLOCK = 6, 6, 384, 256 |
|
|
| @st.cache_resource |
| def load_model(): |
| enc = tiktoken.get_encoding("gpt2") |
| model = build_model(enc.n_vocab, N_LAYER, N_HEAD, N_EMBD, BLOCK).to(DEVICE).eval() |
| state = torch.load(CKPT, map_location=DEVICE) |
| model.load_state_dict(state) |
| return model, enc |
|
|
| st.title("LittleGPT Space ✨") |
| model, enc = load_model() |
|
|
| prompt = st.text_area("Enter your prompt", "Once upon a time...") |
| max_new = st.slider("Max new tokens", 16, 256, 100, 8) |
| temp = st.slider("Temperature", 0.1, 1.5, 0.9, 0.1) |
| top_k = st.slider("Top-k (0 = off)", 0, 200, 50, 10) |
|
|
| if st.button("Generate"): |
| x = torch.tensor([enc.encode(prompt)], dtype=torch.long, device=DEVICE) |
| with torch.no_grad(): |
| y = model.generate(x, max_new_tokens=max_new, temperature=temp, top_k=top_k or None) |
| st.write(enc.decode(y[0].tolist())) |
|
|
