better error handling

app.py

@@ -334,56 +334,73 @@ gpt_rev.eval()
 ##################################################################################
 
 def generate_names(name_start, name_end, number_of_names, model):
+    if number_of_names < 0:
+        return "Error: Please enter a positive number of names to generate!"
+    
+    # Select model
     if model == "MLP":
         config = mlp_config
+        sample_fcn = mlp.sample_char
     elif model == "WaveNet":
         config = wavenet_config
+        sample_fcn = wavenet.sample_char
     elif model == "GPT Micro":
         config = gpt_micro_config
+        sample_fcn = gpt_micro.sample_char
     elif model == "GPT Rev":
         config = gpt_rev_config
+        sample_fcn = gpt_rev.sample_char
+    elif model == "GPT First Rev":
+        # TODO: Change model!
+        config = gpt_rev_config
+        sample_fcn = gpt_rev.sample_char
     else:
-        raise Exception("Model not selected")
+        return "Error: Model not selected"
 
     stoi = config['stoi']
     itos = {s:i for i,s in stoi.items()}
 
-    names = ""
+    output = ""
+
+    # Sanitize user inputs, and append errors to output
+    name_end = name_end.lower()
+    name_start = name_start.lower()
+
+    for c in name_end:
+        if c not in stoi:
+            return "Please change name end. \"" + c + "\" not included in the training set."
+
+    for c in name_start:
+        if c not in stoi:
+            return "Please change name start. \"" + c + "\" not included in the training set."
+
+    if "num_final_chars_in_dataset" in config and len(name_end) > config["num_final_chars_in_dataset"]:
+        name_end = name_end[-config["num_final_chars_in_dataset"]:]
+        output += "Only accepts up to " + str(config["num_final_chars_in_dataset"]) + " final chars. Using: " + str(name_end) + "\n"
+            
+    elif "num_final_chars_in_dataset" not in config and name_end != "":
+        output += "Final chars not used. Need to use a \"Rev\" model trained with this feature.\n"
+
+
+    ## Print requested names
     for _ in range((int)(number_of_names)):
         name = ""
         context = [0] * config['window']
         
         if "num_final_chars_in_dataset" in config:
-            # Put final chars in context
-            if len(name_end) > config["num_final_chars_in_dataset"]:
-                name_end = name_end[-config["num_final_chars_in_dataset"]:]
-                print("Only accepts up to " + str(config["num_final_chars_in_dataset"]) + " final chars. Using: " + name_end)
-            
             for c in name_end:
                 context = context[1:] + [stoi[c]]
             context = context[1:] + [stoi['.']]
-            
-        elif (name_end != ""):
-            print("Final chars not used. Need to use a model trained with this feature.")
     
         # Initialize name with user input
-        for c in name_start.lower():
+        for c in name_start:
             name += c
             context = context[1:] + [stoi[c]]
 
         # Run inference to finish off the name
         while True:
             x = torch.tensor(context).view(1, -1)
-            if model == "MLP":
-                ix = mlp.sample_char(x)
-            elif model == "WaveNet":
-                ix = wavenet.sample_char(x)
-            elif model == "GPT Micro":
-                ix = gpt_micro.sample_char(x)
-            elif model == "GPT Rev":
-                ix = gpt_rev.sample_char(x)
-            else:
-                raise Exception("Model not selected")
+            ix = sample_fcn(x)
                 
             context = context[1:] + [ix]
             name += itos[ix]
@@ -391,9 +408,9 @@ def generate_names(name_start, name_end, number_of_names, model):
             if ix == 0:
                 break
             
-        names += name + "\n"
+        output += name + "\n"
         
-    return names
+    return output
 
 demo = gr.Interface(
     fn=generate_names,
@@ -401,7 +418,7 @@ demo = gr.Interface(
         gr.Textbox(placeholder="Start name with..."),
         gr.Textbox(placeholder="End name with... (only works for rev model)"),
         gr.Number(value=5),
-        gr.Dropdown(["MLP", "WaveNet", "GPT Micro", "GPT Rev"], value="GPT Rev"),
+        gr.Dropdown(["MLP", "WaveNet", "GPT Micro", "GPT Rev", "GPT First Rev"], value="GPT Rev"),
     ],
     outputs="text",
 )