update models

app.py

@@ -7,15 +7,12 @@ from utils import (
     embeddings_predict_relevance,
     stsb_model,
     stsb_tokenizer,
-    ms_model,
-    ms_tokenizer,
     cross_encoder_predict_relevance
 )
 
 def predict(system_prompt, user_prompt):
     predicted_label_jina, probabilities_jina = embeddings_predict_relevance(system_prompt, user_prompt, jina_model, jina_tokenizer, device)
     predicted_label_stsb, probabilities_stsb = cross_encoder_predict_relevance(system_prompt, user_prompt, stsb_model, stsb_tokenizer, device)
-    predicted_label_ms, probabilities_ms = cross_encoder_predict_relevance(system_prompt, user_prompt, ms_model, ms_tokenizer, device)
 
     result = f"""
     **Prediction Summary**
@@ -27,10 +24,6 @@ def predict(system_prompt, user_prompt):
     **2. Model: cross-encoder/stsb-roberta-base**
     - **Prediction**: {"🟥 Off-topic" if predicted_label_stsb==1 else "🟩 On-topic"}
     - **Probability of being off-topic**: {probabilities_stsb[0][1]:.2%}
-
-    **3. Model: cross-encoder/ms-marco-MiniLM-L-6-v2**
-    - **Prediction**: {"🟥 Off-topic" if predicted_label_ms==1 else "🟩 On-topic"}
-    - **Probability of being off-topic**: {probabilities_ms[0][1]:.2%}
     """
 
     return result

models/cross-encoder-ms-marco-MiniLM-L-6-v2-CrossEncoder-OffTopic-Classifier-20240918-090615.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:78a99fac3bc5b4729fee844d2154ea625aa9ceac2928cd648984ee1da5b8a203
-size 91236352

models/cross-encoder-stsb-roberta-base-CrossEncoder-OffTopic-Classifier-20240920-174009.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1e90752828e92bc2f8ec567b85b3de5a0c8c5ddc331c1907d4dfa950624f71ce
-size 500085976

models/jinaai-jina-embeddings-v2-small-en-TwinEncoder-OffTopic-Classifier-20240915-151858.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:223687abc28cf0fa198d326d2786374000396d841e66d684c022941da2ca9628
-size 144076480

utils.py

@@ -1,8 +1,11 @@
+import json
 import torch
 from torch import nn
 from safetensors.torch import load_file
 from transformers import AutoModel, AutoTokenizer
+from huggingface_hub import hf_hub_download
 
+# Set device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Load the model state_dict from safetensors
@@ -13,9 +16,9 @@ def load_model_safetensors(model, load_path="model.safetensors"):
     model.load_state_dict(state_dict)
     return model
 
-##########################
+###################
 # JINA EMBEDDINGS
-##########################
+###################
 
 # Jina Configs
 JINA_CONTEXT_LEN = 1024
@@ -101,7 +104,7 @@ class CrossEncoderWithSharedBase(nn.Module):
         logits = self.classifier(projected)
         return logits
 
-# Prediction function
+# Prediction function for embeddings relevance
 def embeddings_predict_relevance(sentence1, sentence2, model, tokenizer, device):
     model.eval()
     inputs1 = tokenizer(sentence1, return_tensors="pt", truncation=True, padding="max_length", max_length=1024)
@@ -117,23 +120,32 @@ def embeddings_predict_relevance(sentence1, sentence2, model, tokenizer, device)
         predicted_label = torch.argmax(probabilities, dim=1).item()
     return predicted_label, probabilities.cpu().numpy()
 
-# Jina model
-JINA_MODEL_NAME = "jinaai/jina-embeddings-v2-small-en"
+# Load configuration file
+jina_repo_path = "govtech/jina-embeddings-v2-small-en-off-topic"
+jina_config_path = hf_hub_download(repo_id=jina_repo_path, filename="config.json")
+with open(jina_config_path, 'r') as f:
+    jina_config = json.load(f)
+
+# Load Jina model configuration
+JINA_MODEL_NAME = jina_config['classifier']['embedding']['model_name']
+jina_model_weights_fp = jina_config['classifier']['embedding']['model_weights_fp']
+
+# Load tokenizer and model
 jina_tokenizer = AutoTokenizer.from_pretrained(JINA_MODEL_NAME)
 jina_base_model = AutoModel.from_pretrained(JINA_MODEL_NAME)
 jina_model = CrossEncoderWithSharedBase(jina_base_model, num_labels=2)
-jina_model = load_model_safetensors(jina_model, load_path="models/jinaai-jina-embeddings-v2-small-en-TwinEncoder-OffTopic-Classifier-20240915-151858.safetensors")
 
-##########################
+# Load model weights from safetensors
+jina_model_weights_path = hf_hub_download(repo_id=jina_repo_path, filename=jina_model_weights_fp)
+jina_model = load_model_safetensors(jina_model, jina_model_weights_path)
+
+#################
 # CROSS-ENCODER
-##########################
+#################
 
-# STSB Configs
+# STSB Configuration
 STSB_CONTEXT_LEN = 512
 
-# ms-macro Configs
-MS_CONTEXT_LEN = 512
-
 class CrossEncoderWithMLP(nn.Module):
     def __init__(self, base_model, num_labels=2):
         super(CrossEncoderWithMLP, self).__init__()
@@ -162,6 +174,7 @@ class CrossEncoderWithMLP(nn.Module):
         logits = self.classifier(mlp_output)
         return logits
 
+# Prediction function for cross-encoder
 def cross_encoder_predict_relevance(sentence1, sentence2, model, tokenizer, device):
     model.eval()
     # Tokenize the pair of sentences
@@ -187,16 +200,20 @@ def cross_encoder_predict_relevance(sentence1, sentence2, model, tokenizer, devi
         predicted_label = torch.argmax(probabilities, dim=1).item()
     return predicted_label, probabilities.cpu().numpy()
 
-# STSB model
-STSB_MODEL_NAME = "cross-encoder/stsb-roberta-base"
+# Load STSB model configuration
+stsb_repo_path = "govtech/stsb-roberta-base-off-topic"
+stsb_config_path = hf_hub_download(repo_id=stsb_repo_path, filename="config.json")
+with open(stsb_config_path, 'r') as f:
+    stsb_config = json.load(f)
+
+STSB_MODEL_NAME = stsb_config['classifier']['embedding']['model_name']
+stsb_model_weights_fp = stsb_config['classifier']['embedding']['model_weights_fp']
+
+# Load STSB tokenizer and model
 stsb_tokenizer = AutoTokenizer.from_pretrained(STSB_MODEL_NAME)
 stsb_base_model = AutoModel.from_pretrained(STSB_MODEL_NAME)
 stsb_model = CrossEncoderWithMLP(stsb_base_model, num_labels=2)
-stsb_model = load_model_safetensors(stsb_model, load_path="models/cross-encoder-stsb-roberta-base-CrossEncoder-OffTopic-Classifier-20240920-174009.safetensors")
-
-# MS model
-MS_MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2"
-ms_tokenizer = AutoTokenizer.from_pretrained(MS_MODEL_NAME)
-ms_base_model = AutoModel.from_pretrained(MS_MODEL_NAME)
-ms_model = CrossEncoderWithMLP(ms_base_model, num_labels=2)
-ms_model = load_model_safetensors(ms_model, load_path="models/cross-encoder-ms-marco-MiniLM-L-6-v2-CrossEncoder-OffTopic-Classifier-20240918-090615.safetensors")
+
+# Load model weights from safetensors for STSB
+stsb_model_weights_path = hf_hub_download(repo_id=stsb_repo_path, filename=stsb_model_weights_fp)
+stsb_model = load_model_safetensors(stsb_model, stsb_model_weights_path)