from transformers import FlaxAutoModel, AutoTokenizer
import flax.linen as nn
import jax.numpy as jnp
import jax.random as jr
from flax.serialization import from_bytes
import jax

hf_model = FlaxAutoModel.from_pretrained('answerdotai/answerai-colbert-small-v1')
hf_tokenizer = AutoTokenizer.from_pretrained('answerdotai/answerai-colbert-small-v1')
flax_module = hf_model.module

class FlaxVespaColBERTModule(nn.Module):
    def setup(self):
        self.bert = flax_module
        self.linear = nn.Dense(96, use_bias=False)

    def __call__(self, **inputs):
        outputs = self.bert(**inputs).last_hidden_state
        outputs = self.linear(outputs)
        outputs = outputs / jnp.linalg.norm(outputs, axis=-1, keepdims=True)
        return outputs


def get_flax_colbert_model():
    model = FlaxVespaColBERTModule()
    sample_text = hf_tokenizer(
        'Hi, this is a colbert model', return_tensors="np", padding="max_length", truncation=True
    )
    init_params = model.init(jr.PRNGKey(0), **sample_text)

    with open('colbert_flax.msgpack', 'rb') as f:
        serialized = f.read()

    custom_flax_params = from_bytes(init_params, serialized)
    del init_params, sample_text, serialized
    return model, custom_flax_params



class Colbert:
    def __init__(self):
        self.colbert_model, self.colbert_params = get_flax_colbert_model()
        self.forward = jax.jit(self.colbert_model.apply)
        self.tokenizer = hf_tokenizer

    def __call__(self, **inputs):
        return self.forward(self.colbert_params, **inputs)