Create app.py

app.py

@@ -0,0 +1,44 @@
+import gradio as gr
+import torch
+import numpy as np
+from model import Transformer
+from transformers import AutoTokenizer  # pip install transformers
+from utils import (
+    DEVICE,
+    DROPOUT,
+    NUM_EMBED,
+    NUM_HEAD,
+    NUM_LAYER,
+    encode,
+    decode
+)
+
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+vocab_size = tokenizer.vocab_size
+
+# train a new model
+model = Transformer(
+    vocab_size=vocab_size,
+    num_embed=NUM_EMBED,
+    block_size=BLOCK_SIZE,
+    num_heads=NUM_HEAD,
+    num_layers=NUM_LAYER,
+    dropout=DROPOUT
+)
+# load model to GPU if available
+m = model.to(DEVICE)
+# print the number of parameters in the model
+
+m.load_state_dict(torch.load("base_model_1.pth", map_location=torch.device(DEVICE)))
+m.eval()
+
+#print(
+#    "Model with {:.2f}M parameters".format(sum(p.numel() for p in m.parameters()) / 1e6)
+#)
+def model_generate(text, number, TEMPERATURE, TOP_K):
+    context = encode(str(text), tokenizer).unsqueeze(0).to(DEVICE)
+
+    return decode(enc_sec=m.generate(idx=context, max_new_tokens=number, temperature = TEMPERATURE, top_k = TOP_K)[0], tokenizer=tokenizer)
+
+iface = gr.Interface(fn=model_generate, inputs=["text", gr.Slider(10, 1000), gr.Slider(0, 1), gr.Slider(0, 1)], outputs="text")
+iface.launch()