Update app.py

app.py

@@ -5,6 +5,7 @@ import numpy as np
 
 from sentence_transformers import SentenceTransformer
 from sentence_transformers.util import cos_sim
+from sentence_transformers import CrossEncoder
 
 
 codes = """001 - Vehicle Registration (New)
@@ -361,6 +362,8 @@ model_ids = [
     "sentence-transformers/distiluse-base-multilingual-cased-v2",
     "Alibaba-NLP/gte-multilingual-base",
     "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+    "BAAI/bge-reranker-v2-m3",
+    "jinaai/jina-reranker-v2-base-multilingual"
 ]
 # model_id = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
 # model_id = "Alibaba-NLP/gte-multilingual-base"
@@ -368,8 +371,18 @@ model_ids = [
 # model_id = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
 # model_id = "intfloat/multilingual-e5-small"
 # model_id = "sentence-transformers/distiluse-base-multilingual-cased-v2"
+
 model_id = model_ids[-1]
-model = SentenceTransformer(model_id, trust_remote_code=True)
+
+if model_id in model_ids[-2:]:
+    model = CrossEncoder(
+        # "jinaai/jina-reranker-v2-base-multilingual",
+        "BAAI/bge-reranker-v2-m3",
+        automodel_args={"torch_dtype": "auto"},
+        trust_remote_code=True,
+    )
+else:
+    model = SentenceTransformer(model_id, trust_remote_code=True)
 
 # codes_emb = model.encode([x[6:] for x in codes])
 codes_emb = model.encode([x["examples"] for x in examples])#.mean(axis=1)
@@ -497,9 +510,15 @@ def respond(
     plates = [" ".join(x).upper() for i,x in enumerate(matches)]
 
     plate_numbers = ", ".join(plates)
-    
-    text_emb = model.encode(message)
-    scores = cos_sim(codes_emb, text_emb).mean(axis=-1)#[:,0]
+
+    if model.config._name_or_path in model_ids[-2:]:
+        # documents = [v["name"] for v in detail_perhitungan.values()]
+        sentence_pairs = [[message, v["name"]] for v in detail_perhitungan.values()]
+        scores = model.predict(sentence_pairs, convert_to_tensor=True)
+        # scores = [x["score"] for x in model.rank(message, documents)]
+    else:
+        text_emb = model.encode(message)
+        scores = cos_sim(codes_emb, text_emb).mean(axis=-1)#[:,0]
 
     scores_argsort = scores.argsort(descending=True)
     weights = [18,8,7,6,5,4,3,2,1]