Create app.py

app.py

@@ -0,0 +1,65 @@
+# Third party
+import huggingface_hub
+import streamlit as st
+import torch
+
+
+class Vocab:
+    def __init__(self, file_path):
+        with open(file_path) as f:
+            tokens = f.read()
+        self.stoi = {token: i for i, token in enumerate(tokens)}
+        self.itos = {i: token for i, token in enumerate(tokens)}
+
+    def encode(self, s: str) -> list[int]:
+        return [self.stoi[c] for c in s]
+
+    def decode(self, vec: list[int]) -> str:
+        return "".join(self.itos[i] for i in vec)
+
+
+def load_vocab():
+    file_path = huggingface_hub.hf_hub_download(
+        repo_id="szymon-piechowicz-wandb/gpt", filename="vocab.txt"
+    )
+    return Vocab(file_path)
+
+
+def load_model():
+    file_path = huggingface_hub.hf_hub_download(
+        repo_id="szymon-piechowicz-wandb/gpt", filename="model_jit.pt"
+    )
+    return torch.jit.load(file_path)
+
+
+def generate(
+    vocab: Vocab, model, block_size: int, context: str = "\n", num_tokens: int = 1000
+) -> str:
+    """
+    Repeatedly calls the model to generate the next token.
+    """
+    # (T)
+    context = torch.tensor(vocab.encode(context), dtype=torch.long)
+    contexts = context.reshape((1, context.shape[0]))  # (batch_size, T)
+    for _ in range(num_tokens):
+        # (batch_size, block_size)
+        contexts_cropped = contexts[:, -block_size:]
+        logits = model(contexts_cropped)  # (batch_size, block_size, vocab_size)
+        # get the last time step
+        logits = logits[:, -1, :]  # (batch_size, vocab_size)
+        # get probabilities
+        # (batch_size, vocab_size)
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        # sample from the distribution
+        next = torch.multinomial(probs, num_samples=1)  # (batch_size, 1)
+        # append to context
+        contexts = torch.cat((contexts, next), dim=1)  # (batch_size, T + 1)
+    return vocab.decode(contexts[0].tolist())
+
+
+if __name__ == "__main__":
+    vocab = load_vocab()
+    model = load_model()
+    block_size = model.position_embedding.weight.shape[0]
+    model.eval()
+    st.text(generate(vocab=vocab, model=model, block_size=block_size))