app update

app.py

@@ -185,27 +185,51 @@ class BigramLanguageModel(nn.Module):
         return idx
 
 
-# Load the model
-loaded_model = BigramLanguageModel().to(device)  # Initialize an instance of your model
-loaded_model.load_state_dict(torch.load('bigram_language_model.pth', map_location=torch.device('cpu')))
-loaded_model.eval()  # Set the model to evaluation mode
+# Load the shakespeaere model
+shakespeare_model = BigramLanguageModel().to(device)  # Initialize an instance of your model
+shakespeare_model.load_state_dict(torch.load('shakespeaere_language_model.pth', map_location=torch.device('cpu')))
+shakespeare_model.eval()  # Set the model to evaluation mode
 
-def generate_gpt_outputs(prompt=None, max_new_tokens=2000):
+# Load the wikipedia model
+wikipedia_model = BigramLanguageModel().to(device)  # Initialize an instance of your model
+wikipedia_model.load_state_dict(torch.load('wikipedia_language_model.pth', map_location=torch.device('cpu')))
+wikipedia_model.eval()  # Set the model to evaluation mode
+
+
+def generate_shakespeare_outputs(prompt=None, max_new_tokens=2000):
   if prompt:
     context = torch.tensor(encode(prompt), dtype=torch.long, device=device).view(1, -1)
   else:
     context = torch.zeros((1, 1), dtype=torch.long, device=device)
-  text_output = decode(loaded_model.generate(context, max_new_tokens=max_new_tokens)[0].tolist())
+  text_output = decode(shakespeare_model.generate(context, max_new_tokens=max_new_tokens)[0].tolist())
   return text_output
 
+
+def generate_wikipedia_outputs(prompt=None, max_new_tokens=2000):
+  if prompt:
+    context = torch.tensor(encode(prompt), dtype=torch.long, device=device).view(1, -1)
+  else:
+    context = torch.zeros((1, 1), dtype=torch.long, device=device)
+  text_output = decode(wikipedia_model.generate(context, max_new_tokens=max_new_tokens)[0].tolist())
+  return text_output
+    
 import gradio as gr
 
-title = "Mini GPT on Shakespeare text"
-description = "Generate an image with a prompt and apply vibrance loss if you wish to"
+title = "Nano GPT"
+description = "Nano GPT trained on Shakespeare and Wikipedia datasets. It is trained on a very small amount of data to understand how GPT's are trained and built. <a href='https://github.com/karpathy/nanoGPT'>The implementation can be found here </a>""
+
+shakespeare_interface = gr.Interface(generate_shakespeare_outputs,
+                    inputs=[gr.Textbox(label="Enter any prompt ", type="text", value="Once upon a time,"),
+                            gr.Slider(minimum=100, maximum=5000, step=100, value=2000, label="Max new tokens")],
+                    outputs=gr.Textbox(label="Output generated", type="text"))
 
-demo = gr.Interface(generate_gpt_outputs,
+wiki_interface = gr.Interface(generate_wikipedia_outputs,
                     inputs=[gr.Textbox(label="Enter any prompt ", type="text", value="Once upon a time,"),
                             gr.Slider(minimum=100, maximum=5000, step=100, value=2000, label="Max new tokens")],
-                    outputs=gr.Textbox(label="Output generated", type="text"), 
-                    title=title, description=description)
+                    outputs=gr.Textbox(label="Output generated", type="text"))
+
+demo = gr.TabbedInterface([shakespeare_interface, wiki_interface], tab_names=["Shakespeare Data", "Wikipedia Data"], 
+                          title=title, description=description)
+
+
 demo.launch()