Upload folder using huggingface_hub

app.py

@@ -1,23 +1,44 @@
 import os
 import gradio as gr
 import torch
-from sentence_transformers import CrossEncoder
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 # ──────────────────────────────────────────────────────────────────────────────
 # Model
 # ──────────────────────────────────────────────────────────────────────────────
 MODEL_ID = os.environ.get("MODEL_ID", "ALJIACHI/Mizan-Rerank-v2")
+MAX_LENGTH = 512
 
-model = CrossEncoder(
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForSequenceClassification.from_pretrained(
     MODEL_ID,
-    max_length=8192,
     trust_remote_code=True,
-    device="cuda" if torch.cuda.is_available() else "cpu",
+    torch_dtype=torch.float32,
 )
+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():
@@ -27,18 +48,17 @@ def rerank(query: str, documents: str) -> str:
     if not docs:
         return "⚠️ يرجى إدخال مستند واحد على الأقل (Please enter at least one document)"
 
-    ranks = model.rank(query, docs)
+    scores = get_scores(query, docs)
+    ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
 
     lines = []
-    for i, entry in enumerate(ranks, 1):
-        idx = entry["corpus_id"]
-        score = entry["score"]
+    for rank, (idx, score) in enumerate(ranked, 1):
         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"### #{i}  —  Score: {score:.4f}\n"
+            f"### #{rank}  —  Score: {score:.4f}\n"
             f"`{bar}`\n\n"
             f"{preview}\n"
         )

requirements.txt

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