Views
No views yet
1from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
2from peft import PeftModel, PeftConfig
3
4def get_model(peft_model_name):
5 config = PeftConfig.from_pretrained(peft_model_name)
6 base_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(config.base_model_name_or_path)
7 model = PeftModel.from_pretrained(base_model, peft_model_name)
8 model = model.merge_and_unload()
9 model.eval()
10 return model
11
12model = get_model('Tevatron/unified-retriever-v0.1').to('cuda:0')
13processor = AutoProcessor.from_pretrained('Tevatron/unified-retriever-v0.1')
141import torch
2from qwen_vl_utils import process_vision_info
3
4def get_embedding(last_hidden_state: torch.Tensor) -> torch.Tensor:
5 reps = last_hidden_state[:, -1]
6 reps = torch.nn.functional.normalize(reps, p=2, dim=-1)
7 return reps
8
9queries = ["Where can we see Llama?", "What is the LLaMA AI model?"]
10
11query_messages = []
12for query in queries:
13 message = [
14 {
15 'role': 'user',
16 'content': [
17 {'type': 'text', 'text': f'Query: {query}'},
18 ]
19 }
20 ]
21 query_messages.append(message)
22
23query_texts = [
24 processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=False) + "<|endoftext|>"
25 for msg in query_messages
26]
27
28query_image_inputs, query_video_inputs = process_vision_info(query_messages)
29query_inputs = processor(text=query_texts, images=query_image_inputs, videos=query_video_inputs, padding='longest', return_tensors='pt').to('cuda:0')
30
31with torch.no_grad():
32 output = model(**query_inputs, return_dict=True, output_hidden_states=True)
33query_embeddings = get_embedding(output.hidden_states[-1])
34[!NOTE] For encoding the textual documents, the code is the same as the above query encoding, but remove the'Query: 'prefix.
1import requests
2from io import BytesIO
3from PIL import Image
4
5# URLs of the images
6url1 = "https://huggingface.co/Tevatron/dse-phi3-docmatix-v2/resolve/main/animal-llama.png"
7url2 = "https://huggingface.co/Tevatron/dse-phi3-docmatix-v2/resolve/main/meta-llama.png"
8
9response1 = requests.get(url1)
10response2 = requests.get(url2)
11
12doc_image1 = Image.open(BytesIO(response1.content))
13doc_image2 = Image.open(BytesIO(response2.content))
14doc_images = [doc_image1, doc_image2]
15
16doc_messages = []
17for doc in doc_images:
18 message = [
19 {
20 'role': 'user',
21 'content': [
22 {'type': 'text', 'text': ''},
23 {'type': 'image', 'image': doc, 'resized_height': 784, 'resized_width': 784}
24
25 ]
26 }
27 ]
28 doc_messages.append(message)
29doc_texts = [
30 processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=False) + "<|endoftext|>"
31 for msg in doc_messages
32]
33
34doc_image_inputs, doc_video_inputs = process_vision_info(doc_messages)
35doc_inputs = processor(text=doc_texts, images=doc_image_inputs, videos=doc_video_inputs, padding='longest', return_tensors='pt').to('cuda:0')
36
37with torch.no_grad():
38 output = model(**doc_inputs, return_dict=True, output_hidden_states=True)
39
40doc_embeddings = get_embedding(output.hidden_states[-1])1from torch.nn.functional import cosine_similarity
2num_queries = query_embeddings.size(0)
3num_passages = doc_embeddings.size(0)
4
5for i in range(num_queries):
6 query_embedding = query_embeddings[i].unsqueeze(0)
7 similarities = cosine_similarity(query_embedding, doc_embeddings)
8 print(f"Similarities for Query {i+1}: {similarities.cpu().float().numpy()}")
9
10# Similarities for Query 1: [0.3282001 0.17449486]
11# Similarities for Query 2: [0.08133292 0.30867738]