LLM model

app.py

@@ -3,6 +3,7 @@ from huggingface_hub import hf_hub_download
 import json
 import tensorflow as tf
 import numpy as np
+from sentence_transformers import SentenceTransformer
 
 
 # Load models
@@ -14,6 +15,10 @@ with open(model_probs_path) as 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)
 
+st_model = SentenceTransformer("avsolatorio/GIST-large-Embedding-v0")
+llm_model_path = hf_hub_download(repo_id="tbitai/gisty-enron1-spam", filename="gisty-enron1-spam.keras")
+llm_model = tf.keras.models.load_model(llm_model_path)
+
 
 # Utils for Bayes
 
@@ -61,9 +66,14 @@ def predict_bayes(text, intr_threshold, unbiased=False):
 def predict_nn(text):
     return nn_model(np.array([text]))[0][0].numpy()
 
+def predict_llm(text):
+    embedding = st_model.encode(text)
+    return 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):
@@ -71,6 +81,8 @@ def predict(model, input_txt, unbiased, intr_threshold):
         return predict_bayes(input_txt, unbiased=unbiased, intr_threshold=intr_threshold)
     elif model == NN:
         return predict_nn(input_txt)
+    elif model == LLM:
+        return predict_llm(input_txt)
 
 
 # UI
@@ -90,11 +102,13 @@ demo = gr.Interface(
     ],
     outputs=[gr.Number(label="Spam probability")],
     title="Bayes or Spam?",
-    description="Choose your model, and predict if your email is a spam! 📨<br>COMING SOON: LLM models.",
+    description="Choose your model, and predict if your email is a spam! 📨",
     examples=[
-        [BAYES, "Enron actuals for June 26, 2000", False, DEFAULT_INTR_THRESHOLD],
+        [BAYES, enron_email := "Enron actuals for June 26, 2000", False, DEFAULT_INTR_THRESHOLD],
         [BAYES, nerissa_email := "Stop the aging clock\nNerissa", False, DEFAULT_INTR_THRESHOLD],
         [BAYES, nerissa_email, True, DEFAULT_INTR_THRESHOLD],
+        [NN, enron_email, None, None],
+        [LLM, enron_email, None, None],
     ],
     article="This is a demo of the models in the [Bayes or Spam?](https://github.com/tbitai/bayes-or-spam) project.",
 )