Update app.py

app.py

@@ -10,10 +10,12 @@ import tldextract
 from ddgs import DDGS
 from rapidfuzz.fuzz import ratio
 
-from sentence_transformers import SentenceTransformer
-from sklearn.metrics.pairwise import cosine_similarity
+from sentence_transformers import (
+    SentenceTransformer,
+    CrossEncoder,
+)
 
-from FlagEmbedding import FlagReranker
+from sklearn.metrics.pairwise import cosine_similarity
 
 # =====================================================
 # APP
@@ -24,22 +26,22 @@ app = FastAPI(
 )
 
 # =====================================================
-# EMBEDDING MODEL
+# LIGHTWEIGHT EMBEDDING MODEL
 # =====================================================
 
 embedding_model = SentenceTransformer(
-    "BAAI/bge-small-en-v1.5"
+    "sentence-transformers/all-MiniLM-L6-v2"
 )
 
 embedding_model.max_seq_length = 256
 
 # =====================================================
-# RERANKER
+# LIGHTWEIGHT RERANKER
 # =====================================================
 
-reranker = FlagReranker(
-    "BAAI/bge-reranker-base",
-    use_fp16=True,
+reranker = CrossEncoder(
+    "cross-encoder/ms-marco-MiniLM-L-6-v2",
+    max_length=256,
 )
 
 # =====================================================
@@ -310,7 +312,10 @@ def split_paragraphs(text):
 
         lower = paragraph.lower()
 
-        if any(junk in lower for junk in JUNK_PHRASES):
+        if any(
+            junk in lower
+            for junk in JUNK_PHRASES
+        ):
             continue
 
         cleaned.append(paragraph)
@@ -332,9 +337,18 @@ def embedding_rank_paragraphs(
 
     try:
 
-        query_embedding = embedding_model.encode([query])
+        query_embedding = embedding_model.encode(
+            [query],
+            normalize_embeddings=True,
+            show_progress_bar=False,
+        )
 
-        paragraph_embeddings = embedding_model.encode(paragraphs)
+        paragraph_embeddings = embedding_model.encode(
+            paragraphs,
+            batch_size=16,
+            normalize_embeddings=True,
+            show_progress_bar=False,
+        )
 
         similarities = cosine_similarity(
             query_embedding,
@@ -376,7 +390,7 @@ def embedding_rank_paragraphs(
         return paragraphs[:20]
 
 # =====================================================
-# RERANKER
+# RERANK PARAGRAPHS
 # =====================================================
 
 def rerank_paragraphs(query, paragraphs):
@@ -391,7 +405,11 @@ def rerank_paragraphs(query, paragraphs):
             for paragraph in paragraphs
         ]
 
-        scores = reranker.compute_score(pairs)
+        scores = reranker.predict(
+            pairs,
+            batch_size=16,
+            show_progress_bar=False,
+        )
 
         ranked = sorted(
             zip(paragraphs, scores),
@@ -491,7 +509,11 @@ def rerank_facts(query, facts):
             for fact in facts
         ]
 
-        scores = reranker.compute_score(pairs)
+        scores = reranker.predict(
+            pairs,
+            batch_size=16,
+            show_progress_bar=False,
+        )
 
         ranked = sorted(
             zip(facts, scores),