Views
No views yet
1import torch
2
3from transformers import AutoTokenizer, AutoModel
4from tqdm import tqdm
5from more_itertools import chunked
6
7
8residual = True
9residual_factor = 0.5
10
11tokenizer = AutoTokenizer.from_pretrained(
12 "Qwen/Qwen3-Embedding-8B",
13 use_fast=True,
14 padding_side='left',
15)
16
17model = AutoModel.from_pretrained(
18 "SituatedEmbedding/SitEmb-v1.5-Qwen3",
19 torch_dtype=torch.bfloat16,
20 device_map={"": 0},
21)
22
23def _pooling(last_hidden_state, attention_mask, pooling, normalize, input_ids=None, match_idx=None):
24 if pooling in ['cls', 'first']:
25 reps = last_hidden_state[:, 0]
26 elif pooling in ['mean', 'avg', 'average']:
27 masked_hiddens = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
28 reps = masked_hiddens.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
29 elif pooling in ['last', 'eos']:
30 left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
31 if left_padding:
32 reps = last_hidden_state[:, -1]
33 else:
34 sequence_lengths = attention_mask.sum(dim=1) - 1
35 batch_size = last_hidden_state.shape[0]
36 reps = last_hidden_state[torch.arange(batch_size, device=last_hidden_state.device), sequence_lengths]
37 elif pooling == 'ext':
38 if match_idx is None:
39 # default mean
40 masked_hiddens = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
41 reps = masked_hiddens.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
42 else:
43 for k in range(input_ids.shape[0]):
44 sep_index = input_ids[k].tolist().index(match_idx)
45 attention_mask[k][sep_index:] = 0
46 masked_hiddens = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
47 reps = masked_hiddens.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
48 else:
49 raise ValueError(f'unknown pooling method: {pooling}')
50 if normalize:
51 reps = torch.nn.functional.normalize(reps, p=2, dim=-1)
52 return reps
53
54
55def first_eos_token_pooling(
56 last_hidden_states,
57 first_eos_position,
58 normalize,
59):
60 batch_size = last_hidden_states.shape[0]
61 reps = last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), first_eos_position]
62 if normalize:
63 reps = torch.nn.functional.normalize(reps, p=2, dim=-1)
64 return reps
65
66def encode_query(tokenizer, model, pooling, queries, batch_size, normalize, max_length, residual):
67 task = "Given a search query, retrieve relevant chunks from fictions that answer the query"
68 sents = []
69 for query in queries:
70 sents.append(get_detailed_instruct(task, query))
71
72 return encode_passage(tokenizer, model, pooling, sents, batch_size, normalize, max_length)
73
74
75def encode_passage(tokenizer, model, pooling, passages, batch_size, normalize, max_length, residual=False):
76 pas_embs = []
77 pas_embs_residual = []
78 total = len(passages) // batch_size + (1 if len(passages) % batch_size != 0 else 0)
79 with tqdm(total=total) as pbar:
80 for sent_b in chunked(passages, batch_size):
81 batch_dict = tokenizer(sent_b, max_length=max_length, padding=True, truncation=True,
82 return_tensors='pt').to(model.device)
83 if residual:
84 batch_list_dict = tokenizer(sent_b, max_length=max_length, padding=True, truncation=True, )
85 input_ids = batch_list_dict['input_ids']
86 attention_mask = batch_list_dict['attention_mask']
87 max_len = len(input_ids[0])
88 input_starts = [max_len - sum(att) for att in attention_mask]
89 eos_pos = []
90 for ii, it in zip(input_ids, input_starts):
91 pos = ii.index(tokenizer.pad_token_id, it)
92 eos_pos.append(pos)
93 eos_pos = torch.tensor(eos_pos).to(model.device)
94 else:
95 eos_pos = None
96 outputs = model(**batch_dict)
97 pemb_ = _pooling(outputs.last_hidden_state, batch_dict['attention_mask'], pooling, normalize)
98 if residual:
99 remb_ = first_eos_token_pooling(outputs.last_hidden_state, eos_pos, normalize)
100 pas_embs_residual.append(remb_)
101 pas_embs.append(pemb_)
102 pbar.update(1)
103 pas_embs = torch.cat(pas_embs, dim=0)
104 if pas_embs_residual:
105 pas_embs_residual = torch.cat(pas_embs_residual, dim=0)
106 else:
107 pas_embs_residual = None
108 return pas_embs, pas_embs_residual
109
110your_query = "Your Query"
111
112query_hidden, _ = encode_query(
113 tokenizer, model, pooling_type="eos", queries=[your_query],
114 batch_size=8, normalize=True, max_length=8192, residual=residual,
115)
116
117passage_affix = "The context in which the chunk is situated is given below. Please encode the chunk by being aware of the context. Context:\n"
118your_chunk = "Your Chunk"
119your_context = "Your Context"
120
121candidate_hidden, candidate_hidden_residual = encode_passage(
122 tokenizer, model, pooling_type="eos", passages=[f"{your_chunk}<|endoftext|>{passage_affix}{your_context}"],
123 batch_size=4, normalize=True, max_length=8192, residual=residual,
124)
125
126query2candidate = query_hidden @ candidate_hidden.T # [num_queries, num_candidates]
127if candidate_hidden_residual is not None:
128 query2candidate_residual = query_hidden @ candidate_hidden_residual.T
129 if residual_factor == 1.:
130 query2candidate = query2candidate_residual
131 elif residual_factor == 0.:
132 pass
133 else:
134 query2candidate = query2candidate * (1. - residual_factor) + query2candidate_residual * residual_factor
135
136print(query2candidate.tolist())