Add UI interface code

app.py

@@ -1,7 +1,35 @@
 import gradio as gr
+import torch
+from model import BigramLM, encode, decode
 
-def greet():
-    return "Hello world!"
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+model = torch.load('saved_model.pth', map_location= torch.device(device))
+
+def generate_text(context, num_of_tokens, temperature=1.0):
+    if context == None or context == '':
+        idx = torch.zeros((1,1), dtype=torch.long)
+    else:
+        idx = torch.tensor(encode(context), dtype=torch.long).unsqueeze(0)
+    
+    return decode(model.generate(idx, max_new_tokens=num_of_tokens,temperature=temperature)[0].tolist())
+
+
+with gr.Blocks as demo:
+    gr.HTML("<h1 align='center'> Shakespeare Text Generator</h1>")
+
+    context = gr.Textbox(label = "Enter context (optional)")
+
+    with gr.Row():
+        num_of_tokens = gr.Number( label = "Max tokens to generate", value = 100)
+        tmp = gr.Slider(label= "Temperature", minimum = 0.0, maximum = 1.0, value = 1.0 )
+    
+    inputs = [
+        context,
+        num_of_tokens,tmp
+    ]
+    generate_btn = gr.Button(value="Generate")
+    outputs = [gr.Textbox(label= "Generated text")]
+    generate_btn.click(fn = generate_text, inputs= inputs, outputs= outputs)
+    
+demo.launch()

model.py

@@ -206,7 +206,7 @@ class BigramLM(nn.Module):
       # sample from the distribution (pick the best)
       idx_next = torch.multinomial(probs, num_samples=1)
       # GPT like output
-      print(decode(idx_next[0].tolist()), end='')
+      #print(decode(idx_next[0].tolist()), end='')
       # append sampled index to running sequence
       idx = torch.cat((idx, idx_next), dim=1)