How do I increase the output dimension when using Late Chunking ?
#68
by
EthanPhelps
- opened
If I want to apply this model to the Late Chunking which requires token-level embeddings, I can only use vectors of 128 dimensions.
Is it possible to increase the dim of these vectors?
tokenizer = AutoTokenizer.from_pretrained(model_path.as_posix())
chunks, span_annotations = chunk_by_sentences(input_text, tokenizer)
emb_model = AutoModel.from_pretrained(model_path.as_posix(), trust_remote_code=True)
tokens = tokenizer(input_text, return_tensors='pt')
model_out: JinaEmbeddingsV4ModelOutput = emb_model(**tokens, task_label="retrieval", )
emb_vectors = chunked_pooling(model_out.multi_vec_emb, [span_annotations])[0]
Hi, you can do the following to get the last token embeddings:
from transformers import AutoModel
m = AutoModel.from_pretrained('jinaai/jina-embeddings-v4', trust_remote_code=True).to('cuda')
t = m.processor.process_texts(['test string', 'testing'], max_length=8192, prefix='Query') # or prefix='Passage'
batch = {k: v.to('cuda') for k, v in t.items()}
out = m(**batch, task_label="retrieval", output_vlm_last_hidden_states=True)
print(out.vlm_last_hidden_states, out.vlm_last_hidden_states.shape)