python: reranking local model

build.sh

#!/bin/bash
podman build -t local-reranker .

Dockerfile

# Use Python 3.11 base image
FROM python:3.11-slim

# Set a working directory inside the container
WORKDIR /app

# Copy your project files
COPY . /app

# Install system dependencies (optional, needed for PyTorch)
RUN apt-get update && apt-get install -y git wget build-essential && rm -rf /var/lib/apt/lists/*

# Install Python dependencies
RUN pip install --upgrade pip
RUN pip install torch sentence-transformers transformers==4.50.2

# Expose port if you plan to run a local API
EXPOSE 8000

# Command to run your script
CMD ["python", "index.py"]

index.py

from sentence_transformers import CrossEncoder

model = CrossEncoder(
#    "BAAI/bge-reranker-base",
    "mixedbread-ai/mxbai-rerank-large-v1",
    device="cpu"  # or "cpu"
)

query = "What is the capital of France?"
docs = ["Paris is the capital of France.",
        "Berlin is the capital of Germany.",
        "Madrid is in Spain."]

pairs = [(query, doc) for doc in docs]
scores = model.predict(pairs)

ranked = sorted(zip(docs, scores), key=lambda x: x[1], reverse=True)
for doc, score in ranked:
    print(score, doc)

run.sh

#!/bin/bash
mkdir -p cache
podman run --rm -it \
  -v ./cache:/root/.cache/huggingface \
  -v $(pwd):/app \
  local-reranker python index.py

result:

bash run.sh
0.99873215 Paris is the capital of France.
0.0084372815 Berlin is the capital of Germany.
0.0034654106 Madrid is in Spain.