# Embedding and similarity using llm model
#
# https://huggingface.co/Qwen/Qwen3-Reranker-0.6B
# Requires transformers>=4.51.0
#
import torch
import numpy as np

from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM
from transformers import set_seed

import os
import logging
import time

from flask import Flask, abort, jsonify, request

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

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


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

def process_inputs(pairs):
    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, **kwargs):
    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)
    scores = batch_scores[:, 1].exp().tolist()
    return scores

def rescore(query, docs):
    with torch.no_grad():
        pairs = [format_instruction(task, query, doc) for doc in docs]
        # Tokenize the input texts
        inputs = process_inputs(pairs)
        scores = compute_logits(inputs)
#    print("scores: ", scores)
    logger.info("scores: " + str(scores))
#    torch.cuda.empty_cache()
    return scores

def empty_cache():
    global cache_cnt
    cache_cnt += 1
    if cache_cnt % 100 == 0:
        torch.cuda.empty_cache()

cache_cnt = 0;

class Args:
    model_url_or_path = 'Qwen/Qwen3-Reranker-0.6B'
#    model_url_or_path = '/app/Qwen3-Reranker-4B'
#    model_url_or_path = '/app/Qwen3-Reranker-0.6B'
args = Args()
logger.info(args)

model_name = args.model_url_or_path
logger.info("Loading model: " + model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
#model = AutoModelForCausalLM.from_pretrained(model_name).eval()
#model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, attn_implementation="flash_attention_2").cuda().eval()
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    dtype="auto",
#    attn_implementation="flash_attention_2",
    device_map="auto"
).eval()
logger.info("Model loaded: " + model_name)

token_false_id = tokenizer.convert_tokens_to_ids("no")
token_true_id = tokenizer.convert_tokens_to_ids("yes")
#max_length = 8192
max_length = 32000

prefix = "<|im_start|>system\nJudge 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"
prefix_tokens = tokenizer.encode(prefix, add_special_tokens=False)
suffix_tokens = tokenizer.encode(suffix, add_special_tokens=False)
        
task = 'Given a web search query, retrieve relevant passages that answer the query'

#====================================
#POST: http://localhost:5007/rescore

app = Flask(__name__)
app.config['CORS_HEADERS'] = 'Content-Type'
app.config['DEBUG'] = False

@app.route("/rescore", methods=['POST'])
def post_rescore():
    data = request.get_json()
    if 'query' not in data or len(data['query']) == 0 or 'docs' not in data or len(data['docs']) == 0:
        abort(500)
    else:
        empty_cache()
        start = time.time()
        query = data['query']
        docs = data['docs']
#        print(query)
        result = rescore(query, docs)
        end = time.time() 
        took = int((end-start)*1000)
        logger.info('Rescoring for %d docs took %d (ms)' % (len(docs), took))
        return jsonify({'result': result,'took': took})

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