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