# Reranking service using Qwen3-Reranker model
#
# https://huggingface.co/Qwen/Qwen3-Reranker-0.6B
# Requires transformers>=4.51.0

import logging
import os
import threading
import time

import torch
from flask import Flask, abort, jsonify, request
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)

seed_value = 1
set_seed(seed_value)
torch.manual_seed(seed_value)
if torch.cuda.is_available():
    torch.cuda.manual_seed(seed_value)
    torch.cuda.manual_seed_all(seed_value)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

MAX_LENGTH = 32000

RERANK_INSTRUCTION = "Given a web search query, retrieve relevant passages that answer the query"

PREFIX = (
    "<|im_start|>system\n"
    "Judge whether the Document meets the requirements based on the Query and the Instruct provided. "
    "Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n"
    "<|im_start|>user\n"
)
SUFFIX = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"

# Serialize model inference — not thread-safe
lock = threading.Lock()

# ---------------------------------------------------------------------------
# Cache helper
# ---------------------------------------------------------------------------

cache_lock = threading.Lock()
cache_cnt = 0

def empty_cache():
    global cache_cnt
    with cache_lock:
        cache_cnt += 1
        flush = (cache_cnt % 10 == 0)
    if flush and torch.cuda.is_available():
        torch.cuda.empty_cache()

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------

MODEL_NAME = os.environ.get("RER_MODEL", "Qwen/Qwen3-Reranker-0.6B")

logger.info("Loading model: %s", MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side='left')
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype="auto",
    device_map="auto",
).eval()
logger.info("Model loaded: %s", MODEL_NAME)

token_false_id = tokenizer.convert_tokens_to_ids("no")
token_true_id  = tokenizer.convert_tokens_to_ids("yes")
prefix_tokens  = tokenizer.encode(PREFIX, add_special_tokens=False)
suffix_tokens  = tokenizer.encode(SUFFIX, add_special_tokens=False)

# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def format_instruction(query: str, doc: str) -> str:
    return "<Instruct>: {instruction}\n<Query>: {query}\n<Document>: {doc}".format(
        instruction=RERANK_INSTRUCTION, query=query, doc=doc
    )


def process_inputs(pairs: list):
    inputs = tokenizer(
        pairs,
        padding=False,
        truncation='longest_first',
        return_attention_mask=False,
        max_length=MAX_LENGTH - len(prefix_tokens) - len(suffix_tokens),
    )
    for i, ele in enumerate(inputs['input_ids']):
        inputs['input_ids'][i] = prefix_tokens + ele + suffix_tokens
    inputs = tokenizer.pad(inputs, padding=True, return_tensors="pt", max_length=MAX_LENGTH)
    for key in inputs:
        inputs[key] = inputs[key].to(model.device)
    return inputs


def compute_logits(inputs) -> list:
    batch_scores = model(**inputs).logits[:, -1, :]
    true_vector  = batch_scores[:, token_true_id]
    false_vector = batch_scores[:, token_false_id]
    batch_scores = torch.stack([false_vector, true_vector], dim=1)
    batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1)
    return batch_scores[:, 1].exp().tolist()


def rescore(query: str, docs: list) -> list:
    pairs = [format_instruction(query, doc) for doc in docs]
    inputs = process_inputs(pairs)
    with lock, torch.no_grad():
        return compute_logits(inputs)

# ---------------------------------------------------------------------------
# Flask app
# ---------------------------------------------------------------------------

# POST: http://localhost:8088/rescore

app = Flask(__name__)
app.config['DEBUG'] = False


@app.route("/rescore", methods=['POST'])
def post_rescore():
    data = request.get_json()
    if not data or not data.get('query') or not data.get('docs'):
        abort(400, description="Request must include non-empty 'query' and 'docs' fields")
    empty_cache()
    query = data['query']
    docs  = data['docs']
    start = time.perf_counter()
    result = rescore(query, docs)
    took = int((time.perf_counter() - start) * 1000)
    logger.info("Rescoring %d docs took %d ms | scores: %s", len(docs), took, result)
    return jsonify({'result': result, 'took': took})


@app.route("/health", methods=['GET'])
def health():
    return jsonify({"status": "ok", "model": MODEL_NAME})


if __name__ == "__main__":
    app.run(host='0.0.0.0', debug=False)
