adds reverse model and feature that allows you to pick last letters in the name

app.py

@@ -34,6 +34,14 @@ gpt_micro_weights_path = huggingface_hub.hf_hub_download(
     "jefsnacker/surname_generator",
     "micro_gpt_weights.pt")
 
+gpt_rev_config_path = huggingface_hub.hf_hub_download(
+    "jefsnacker/surname_generator",
+    "rev_gpt_config.yaml")
+
+gpt_rev_weights_path = huggingface_hub.hf_hub_download(
+    "jefsnacker/surname_generator",
+    "rev_gpt_weights.pt")
+
 with open(mlp_config_path, 'r') as file:
     mlp_config = yaml.safe_load(file)
 
@@ -43,6 +51,9 @@ with open(wavenet_config_path, 'r') as file:
 with open(gpt_micro_config_path, 'r') as file:
     gpt_micro_config = yaml.safe_load(file)
 
+with open(gpt_rev_config_path, 'r') as file:
+    gpt_rev_config = yaml.safe_load(file)
+
 ##################################################################################
 ## MLP
 ##################################################################################
@@ -314,31 +325,48 @@ gpt_micro = GPT(gpt_micro_config)
 gpt_micro.load_state_dict(torch.load(gpt_micro_weights_path))
 gpt_micro.eval()
 
+gpt_rev = GPT(gpt_rev_config)
+gpt_rev.load_state_dict(torch.load(gpt_rev_weights_path))
+gpt_rev.eval()
+
 ##################################################################################
 ## Gradio App
 ##################################################################################
 
-def generate_names(name_start, number_of_names, model):
+def generate_names(name_start, name_end, number_of_names, model):
     if model == "MLP":
-        stoi = mlp_config['stoi']
-        window = mlp_config['window']
+        config = mlp_config
     elif model == "WaveNet":
-        stoi = wavenet_config['stoi']
-        window = wavenet_config['window']
+        config = wavenet_config
     elif model == "GPT Micro":
-        stoi = gpt_micro_config['stoi']
-        window = gpt_micro_config['window']
+        config = gpt_micro_config
+    elif model == "GPT Rev":
+        config = gpt_rev_config
     else:
         raise Exception("Model not selected")
-        
+
+    stoi = config['stoi']
     itos = {s:i for i,s in stoi.items()}
 
     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
-        name = ""
-        context = [0] * window
         for c in name_start.lower():
             name += c
             context = context[1:] + [stoi[c]]
@@ -352,6 +380,8 @@ def generate_names(name_start, number_of_names, model):
                 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")
                 
@@ -369,8 +399,9 @@ demo = gr.Interface(
     fn=generate_names,
     inputs=[
         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"], value="GPT Micro"),
+        gr.Dropdown(["MLP", "WaveNet", "GPT Micro", "GPT Rev"], value="GPT Rev"),
     ],
     outputs="text",
 )