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