Separate model names from display names

app.py

@@ -7,19 +7,26 @@ import numpy as np
 
 # Load models
 
-model_probs_path = hf_hub_download(repo_id="tbitai/bayes-enron1-spam", filename="probs.json")
+MODELS = [
+    (BAYES := "bayes-enron1-spam", "Bayes Enron1 spam"),
+    (NN := "nn-enron1-spam", "NN Enron1 spam"),
+    (LLM := "gisty-enron1-spam", "GISTy Enron1 spam"),
+]
+
+model_probs_path = hf_hub_download(repo_id=f"tbitai/{BAYES}", filename="probs.json")
 with open(model_probs_path) as f:
     model_probs = json.load(f)
 
-nn_model_path = hf_hub_download(repo_id="tbitai/nn-enron1-spam", filename="nn-enron1-spam.keras")
+nn_model_path = hf_hub_download(repo_id=f"tbitai/{NN}", filename="nn-enron1-spam.keras")
 nn_model = tf.keras.models.load_model(nn_model_path)
 
-llm_model_path = hf_hub_download(repo_id="tbitai/gisty-enron1-spam", filename="gisty-enron1-spam.keras")
+llm_model_path = hf_hub_download(repo_id=f"tbitai/{LLM}", filename="gisty-enron1-spam.keras")
 llm_model = tf.keras.models.load_model(llm_model_path)
 # Sentence Transformers should be imported after Keras models, in order to prevent it from setting Keras to legacy.
 from sentence_transformers import SentenceTransformer
 st_model = SentenceTransformer("avsolatorio/GIST-large-Embedding-v0")
 
+
 # Utils for Bayes
 
 UNK = '[UNK]'
@@ -70,12 +77,6 @@ def predict_llm(text):
     embedding = st_model.encode(text)
     return float(llm_model(np.array([embedding]))[0][0].numpy())
 
-MODELS = [
-    BAYES := "Bayes Enron1 spam",
-    NN := "NN Enron1 spam",
-    LLM := "GISTy Enron1 spam",
-]
-
 def predict(model, input_txt, unbiased, intr_threshold):
     if model == BAYES:
         return predict_bayes(input_txt, unbiased=unbiased, intr_threshold=intr_threshold)