NN model

app.py

@@ -4,10 +4,19 @@ import json
 import tensorflow as tf
 import numpy as np
 
+
+# Load models
+
 model_probs_path = hf_hub_download(repo_id="tbitai/bayes-enron1-spam", 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 = tf.keras.models.load_model(nn_model_path)
+
+
+# Utils for Bayes
+
 UNK = '[UNK]'
 
 def tokenize(text):
@@ -27,9 +36,14 @@ def get_interesting_probs(probs, intr_threshold):
                   key=lambda p: abs(p - 0.5),
                   reverse=True)[:intr_threshold]
 
+DEFAULT_INTR_THRESHOLD = 15
+
 def unbias(p):
     return (2 * p) / (p + 1)
 
+
+# Predict functions
+
 def predict_bayes(text, intr_threshold, unbiased=False):
     words = tokenize(text)
     probs = []
@@ -44,29 +58,39 @@ def predict_bayes(text, intr_threshold, unbiased=False):
     interesting_probs = get_interesting_probs(probs, intr_threshold)
     return combine(interesting_probs)
 
-DEFAULT_INTR_THRESHOLD = 15
+def predict_nn(text):
+    return nn_model(np.array([text]))[0][0].numpy()
 
 MODELS = [
     BAYES := "Bayes Enron1 spam",
+    NN := "NN Enron1 spam",
 ]
 
 def predict(model, unbiased, intr_threshold, input_txt):
     if model == BAYES:
         return predict_bayes(input_txt, unbiased=unbiased, intr_threshold=intr_threshold)
+    elif model == NN:
+        return predict_nn(input_txt)
+
+
+# UI
 
 demo = gr.Interface(
     fn=predict,
     inputs=[
         gr.Dropdown(choices=MODELS, value=BAYES, label="Model"),
+        gr.TextArea(label="Email"),
+    ],
+    additional_inputs_accordion="Additional configuration for Bayes",
+    additional_inputs=[
         gr.Checkbox(label="Unbias", info="Correct Graham's bias?"),
         gr.Slider(minimum=1, maximum=DEFAULT_INTR_THRESHOLD + 5, step=1, value=DEFAULT_INTR_THRESHOLD, 
                   label="Interestingness threshold", 
                   info=f"How many of the most interesting words to select in the probability calculation? ({DEFAULT_INTR_THRESHOLD} for Graham)"),
-        gr.TextArea(label="Email"),
     ],
     outputs=[gr.Number(label="Spam probability")],
     title="Bayes or Spam?",
-    description="Choose and configure your model, and predict if your email is a spam! 📨<br>COMING SOON: NN and LLM models.",
+    description="Choose your model, and predict if your email is a spam! 📨<br>COMING SOON: LLM models.",
     examples=[
         [BAYES, False, DEFAULT_INTR_THRESHOLD, "Enron actuals for June 26, 2000"],
         [BAYES, False, DEFAULT_INTR_THRESHOLD, nerissa_email := "Stop the aging clock\nNerissa"],