# Embedding and similarity using Qwen3-Embedding-0.6B model
#
# https://huggingface.co/Qwen/Qwen3-Embedding-0.6B
# Requires transformers>=4.51.0, sentence-transformers>=2.7.0
#
# Instruction-aware: queries use prompt_name="query" (adds task instruction prefix)
# Documents are encoded as plain text (no instruction needed)
# Native 1024 dimensions, 32k context length
#
import torch
import numpy as np

from sentence_transformers import SentenceTransformer
from transformers import set_seed

import os
import logging
import time

from flask import Flask, abort, jsonify, request

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__)

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
    logger.info("Cuda available!")

def round_floats(o):
    if isinstance(o, float): return round(o, 6)
    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 float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

def embed_query(query):
    # prompt_name="query" applies the instruction prefix automatically
    emb = model.encode([query], prompt_name="query", normalize_embeddings=True)
    return round_floats(emb.tolist()[0])

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

def embedding(texts):
    embs = model.encode(texts, normalize_embeddings=True, batch_size=32)
    return embs.tolist()

def similarity(text, texts):
    query_embedding = model.encode([text], prompt_name="query", normalize_embeddings=True)[0]
    doc_embeddings = model.encode(texts, normalize_embeddings=True, batch_size=32)
    return [cosine_similarity(query_embedding, doc_emb) for doc_emb in doc_embeddings]

cache_cnt = 0

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

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

model_name = args.model_url_or_path
logger.info("Loading model: " + model_name)

try:
    import flash_attn
    attn_implementation = "flash_attention_2"
    logger.info("Flash Attention 2 available, using it")
except ImportError:
    attn_implementation = "eager"
    logger.info("Flash Attention 2 not available, using default attention")

model = SentenceTransformer(
    model_name,
    model_kwargs={
        "attn_implementation": attn_implementation,
        "device_map": "auto",
        "torch_dtype": torch.bfloat16,
    },
    tokenizer_kwargs={"padding_side": "left"},
)
logger.info("Model loaded: " + model_name)


#====================================
#http://localhost:8085/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 not text or len(text) == 0:
        abort(500)
    else:
        empty_cache()
        start = time.time()
        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': round_floats(result), 'took': took})

@app.route("/embed", methods=['POST'])
def post_embed():
    data = request.get_json()
    isQuery = data.get('isQuery', False)
    if 'text' not in data or len(data['text']) == 0:
        abort(500)
    else:
        empty_cache()
        start = time.time()
        text = data['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': round_floats(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']
        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']
        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)