Upload folder using huggingface_hub

app.py

@@ -1,44 +1,23 @@
 import os
 import gradio as gr
 import torch
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from sentence_transformers import CrossEncoder
 
 # ──────────────────────────────────────────────────────────────────────────────
 # Model
 # ──────────────────────────────────────────────────────────────────────────────
 MODEL_ID = os.environ.get("MODEL_ID", "ALJIACHI/Mizan-Rerank-v2")
-MAX_LENGTH = 512
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForSequenceClassification.from_pretrained(
+model = CrossEncoder(
     MODEL_ID,
+    max_length=8192,
     trust_remote_code=True,
-    torch_dtype=torch.float32,
+    device="cuda" if torch.cuda.is_available() else "cpu",
 )
-model.eval()
 
 # ──────────────────────────────────────────────────────────────────────────────
 # Reranking logic
 # ──────────────────────────────────────────────────────────────────────────────
-def get_scores(query: str, passages: list[str]) -> list[float]:
-    """Score query-passage pairs using the model directly."""
-    pairs = [(query, p) for p in passages]
-    inputs = tokenizer(
-        [q for q, _ in pairs],
-        [p for _, p in pairs],
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        max_length=MAX_LENGTH,
-    )
-    with torch.no_grad():
-        outputs = model(**inputs)
-    scores = torch.sigmoid(outputs.logits.squeeze(-1)).tolist()
-    if isinstance(scores, float):
-        scores = [scores]
-    return scores
-
-
 def rerank(query: str, documents: str) -> str:
     """Score and rerank documents against a query."""
     if not query.strip():
@@ -48,17 +27,18 @@ def rerank(query: str, documents: str) -> str:
     if not docs:
         return "⚠️ يرجى إدخال مستند واحد على الأقل (Please enter at least one document)"
 
-    scores = get_scores(query, docs)
-    ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
+    ranks = model.rank(query, docs)
 
     lines = []
-    for rank, (idx, score) in enumerate(ranked, 1):
+    for i, entry in enumerate(ranks, 1):
+        idx = entry["corpus_id"]
+        score = entry["score"]
         doc_text = docs[idx]
         preview = doc_text[:200] + ("..." if len(doc_text) > 200 else "")
         bar_len = int(score * 20)
         bar = "█" * bar_len + "░" * (20 - bar_len)
         lines.append(
-            f"### #{rank}  —  Score: {score:.4f}\n"
+            f"### #{i}  —  Score: {score:.4f}\n"
             f"`{bar}`\n\n"
             f"{preview}\n"
         )

requirements.txt

@@ -1,3 +1,4 @@
+sentence-transformers>=5.0.0
 transformers>=4.40.0
 torch>=2.0.0
 gradio>=5.0.0