Update app.py

app.py

@@ -7,15 +7,25 @@ from utils import GPTLM
 
 newmodel = GPTLM.load_from_checkpoint('shakespeare_gpt.pth')
 
-def generate_art(character_dropdown, seed_slider):
+chars = ['\n', ' ', '!', '$', '&', "'", ',', '-', '.', '3', ':', ';', '?', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
+vocab_size = len(chars)
+# create a mapping from characters to integers
+stoi = { ch:i for i,ch in enumerate(chars) }
+itos = { i:ch for i,ch in enumerate(chars) }
 
-  if character_dropdown == "NONE":
-    return "NULL"
-  else:
-    return "NOT NULL"
+encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers
+decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string
 
-  return "Hello There!"
 
+def generate_dialogue(character_dropdown, seed_slider):
+    
+  if character_dropdown == "NONE":
+    context = torch.zeros((1, 1), dtype=torch.long)
+    return decode(newmodel.model.generate(context, max_new_tokens=100)[0].tolist())
+  else:
+    context = torch.tensor([encode(character_dropdown)], dtype=torch.long)
+    return decode(newmodel.model.generate(context, max_new_tokens=100)[0].tolist())
+      
 
 HTML_TEMPLATE = """    
 <style>
@@ -116,7 +126,7 @@ with gr.Blocks() as interface:
 
     with gr.Row():
         button = gr.Button("Generate Dialogue")
-        button.click(generate_art, inputs=inputs, outputs=outputs)
+        button.click(generate_dialogue, inputs=inputs, outputs=outputs)
 
 
 if __name__ == "__main__":