Views
No views yet
1import os
2os.environ["CUDA_VISIBLE_DEVICES"] = "0"
3
4import torch
5from torch import nn
6from llm2vec import LLM2Vec
7from transformers import AutoTokenizer, AutoModel, AutoConfig
8
9class LLM2VecWrapper(LLM2Vec):
10 def prepare_for_tokenization(self, text):
11 text = (
12 "<|start_header_id|>user<|end_header_id|>\n\n"
13 + text.strip()
14 + "<|eot_id|>"
15 )
16 return text
17
18class LlamaVec_1B_FeatureExtractor(nn.Module):
19 def __init__(self):
20 super().__init__()
21
22 model_path = 'microsoft/LLM2CLIP-Llama-3.2-1B-Instruct-CC-Finetuned'
23 config = AutoConfig.from_pretrained(model_path)
24
25 model = AutoModel.from_pretrained(model_path, config=config, trust_remote_code=True, torch_dtype=torch.bfloat16)
26 tokenizer = AutoTokenizer.from_pretrained(model_path)
27
28 tokenizer.pad_token = tokenizer.eos_token
29 tokenizer.padding_side = "left"
30
31 self.l2v = LLM2VecWrapper(model, tokenizer, pooling_mode="mean", max_length=512, skip_instruction=True)
32
33 def extract_features(self, text):
34 with torch.amp.autocast('cuda'):
35 reps_norm = self.l2v.encode(text)
36 reps_norm = torch.nn.functional.normalize(reps_norm, p=2, dim=1)
37 return {"embeds": reps_norm}
38
39text_model = LlamaVec_1B_FeatureExtractor()
40captions = ["this is a test"]
41embeddings = text_model.extract_features(captions)@misc{huang2024llm2clippowerfullanguagemodel,
title={LLM2CLIP: Powerful Language Model Unlock Richer Visual Representation},
author={Weiquan Huang and Aoqi Wu and Yifan Yang and Xufang Luo and Yuqing Yang and Liang Hu and Qi Dai and Xiyang Dai and Dongdong Chen and Chong Luo and Lili Qiu},
year={2024},
eprint={2411.04997},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2411.04997},
}