# Embedding and similarity using Qwen3 model
#
# https://github.com/QwenLM/Qwen3-Embedding
# Requires transformers>=4.51.0
#
import os
import logging
import time

import numpy as np

import torch
from torch import Tensor
#from torch import nn
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel
from transformers.utils import is_flash_attn_2_available
from typing import Dict, Optional, List, Union

from transformers import set_seed

from flask import Flask, abort, jsonify, request

seed_value = 1
set_seed(seed_value)
torch.manual_seed(seed_value)
use_cuda = False
if torch.cuda.is_available():
    use_cuda = True
    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 round_byte(a):
    b = round(a * 740) # for 1024 size
    if b > 127: b = 127
    elif b < -128: b = -128
    return b

def round_floats(o):
    if isinstance(o, float): return round_byte(o)
    if isinstance(o, dict): return {k: round_floats(v) for k, v in o.items()}
    if isinstance(o, (list, tuple)): return [round_floats(x) for x in o]
    return o

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

def get_detailed_instruct(task_description: str, query: str) -> str:
    return f'Instruct: {task_description}\nQuery: {query}'

def embed_query(query):
    emb = model.encode([query], is_query=True, dim=dim)
    remb = round_floats(emb[0].tolist())
    return remb

def embed(text):
    emb = model.encode([text], dim=dim)
    remb = round_floats(emb[0].tolist())
    return remb

def embedding(texts):
    embs = model.encode(texts, dim=dim)
    rembs = round_floats(embs.tolist())
    return rembs

def similarity(text, texts):
    query_outputs = model.encode([text], is_query=True, dim=dim)
    doc_outputs = model.encode(texts, dim=dim)
    scores = (query_outputs @ doc_outputs.T)[0].tolist()
    logger.info("scores:" + str(scores))
    return scores

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

class Qwen3Embedding():
    def __init__(self, model_name_or_path, instruction=None,  use_fp16: bool = True, use_cuda: bool = True, max_length=512): # max_length=8192
        if instruction is None:
            instruction = 'Given a web search query, retrieve relevant passages that answer the query'
        self.instruction = instruction
        if is_flash_attn_2_available() and use_cuda:
            self.model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True, attn_implementation="flash_attention_2", dtype=torch.float16)
        else:
            self.model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True, dtype=torch.float16)
        if use_cuda:
            self.model = self.model.cuda()
        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True, padding_side='left')
        self.max_length=max_length
    
    def last_token_pool(self, last_hidden_states: Tensor,
        attention_mask: Tensor) -> Tensor:
        left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
        if left_padding:
            return last_hidden_states[:, -1]
        else:
            sequence_lengths = attention_mask.sum(dim=1) - 1
            batch_size = last_hidden_states.shape[0]
        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]

    def get_detailed_instruct(self, task_description: str, query: str) -> str:
        if task_description is None:
            task_description = self.instruction
        return f'Instruct: {task_description}\nQuery:{query}'

    def encode(self, sentences: Union[List[str], str], is_query: bool = False, instruction=None, dim: int = -1):
        if isinstance(sentences, str):
            sentences = [sentences]
        if is_query:
            sentences = [self.get_detailed_instruct(instruction, sent) for sent in sentences]
        inputs = self.tokenizer(sentences, padding=True, truncation=True, max_length=self.max_length, return_tensors='pt')
        inputs.to(self.model.device)
        with torch.no_grad():
            model_outputs = self.model(**inputs)
            output = self.last_token_pool(model_outputs.last_hidden_state, inputs['attention_mask'])
            if dim != -1:
                output = output[:, :dim]
            output  = F.normalize(output, p=2, dim=1)
        return output


class Args:
    model_url_or_path = 'Qwen/Qwen3-Embedding-0.6B'
#    model_url_or_path = '/app/Qwen3-Embedding-0.6B'

args = Args()
logger.info(args)

model_name = args.model_url_or_path
logger.info("Loading model: " + model_name)
model = Qwen3Embedding(model_name, use_cuda = use_cuda)
logger.info("Model loaded: " + model_name)
dim = 1024

#====================================
#http://localhost:5006/embed?text=test

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

@app.route("/embed", methods=['GET'])
def get_embed():
    isQuery = request.args.get('isQuery')
    text = request.args.get('text')
    if len(text) == 0:
        abort(500)
    else:
        empty_cache()
        start = time.time()
#        print(text)
        if isQuery == 'true':
            result = embed_query(text)
        else:   
            result = embed(text)
        end = time.time()
        took = int((end-start)*1000)
        logger.info('Embed took %d (ms)' % (took))
        return jsonify({'result': result,'took': took})

@app.route("/embed", methods=['POST'])
def post_embed():
    data = request.get_json()
    if 'isQuery' in data:
        isQuery = data['isQuery']
    else:
        isQuery = False
    if 'text' not in data or len(data['text']) == 0:
        abort(500)
    else:
        empty_cache()
        start = time.time() 
        text = data['text']
#        print(text)
        if isQuery:
            result = embed_query(text)
        else:
            result = embed(text)
        end = time.time() 
        took = int((end-start)*1000)
        logger.info('Embed took %d (ms)' % (took))
        return jsonify({'result': result,'took': took})

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

@app.route("/similarity", methods=['POST'])
def post_similarity():
    data = request.get_json()
    if 'text' not in data or len(data['text']) == 0 or 'texts' not in data or len(data['texts']) == 0:
        abort(500)
    else:
        empty_cache()
        start = time.time()
        text = data['text']
        texts = data['texts']
#        print(text)
        result = similarity(text, texts)
        end = time.time() 
        took = int((end-start)*1000)
        logger.info('Similarity for %d texts took %d (ms)' % (len(texts), 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)
