import open_clip
import torch
import functools
import onnxruntime as ort

model_name = 'ViT-SO400M-16-SigLIP2-512'
tokenizer = open_clip.get_tokenizer(model_name)
#onnx_path = "/home/thebears/Source/ml/clip_extract/text/ViT-SO400M-16-SigLIP2-512_text"
onnx_path = "/home/thebears/Source/ml/clip_extract/text/ViT-SO400M-16-SigLIP2-512_text_optimized"
ort_sess = ort.InferenceSession(onnx_path)

from bottle import route, run, template, request, debug
@route('/encode')
def get_matches():
    query = request.query.get('query',None)
    if query is None:
        return None
    return execute_model(query)

@functools.cache
def execute_model(query):
    with torch.no_grad():
        text_tokenized = tokenizer(query)
        vec = ort_sess.run(None, {'tokenized': text_tokenized.numpy() })[0].tolist()
    return {'vector':vec}

debug(True)
run(host='0.0.0.0', port=53004, server='bjoern')