Generate late-chunked embeddings given character-offset spans.
| 19 | import numpy as np |
| 20 | |
| 21 | class LateChunkEncoder: |
| 22 | """Generate late-chunked embeddings given character-offset spans.""" |
| 23 | |
| 24 | def __init__(self, model_name: str = "Qwen/Qwen3-Embedding-0.6B", *, max_tokens: int = 8192) -> None: |
| 25 | self.model_name = model_name |
| 26 | self.max_len = max_tokens |
| 27 | self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 28 | # Back-compat: allow short alias without repo namespace |
| 29 | repo_id = model_name |
| 30 | if "/" not in model_name and not model_name.startswith("Qwen/"): |
| 31 | # map common alias to official repo |
| 32 | alias_map = { |
| 33 | "qwen3-embedding-0.6b": "Qwen/Qwen3-Embedding-0.6B", |
| 34 | } |
| 35 | repo_id = alias_map.get(model_name.lower(), model_name) |
| 36 | |
| 37 | self.tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True) |
| 38 | self.model = AutoModel.from_pretrained(repo_id, trust_remote_code=True) |
| 39 | self.model.to(self.device) |
| 40 | self.model.eval() |
| 41 | |
| 42 | @torch.inference_mode() |
| 43 | def encode(self, text: str, chunk_spans: List[Tuple[int, int]]) -> List[np.ndarray]: |
| 44 | """Return one vector *per* span. |
| 45 | |
| 46 | Args: |
| 47 | text: Full document text. |
| 48 | chunk_spans: List of (char_start, char_end) offsets for each chunk. |
| 49 | |
| 50 | Returns: |
| 51 | List of numpy float32 arrays – one per chunk. |
| 52 | """ |
| 53 | if not chunk_spans: |
| 54 | return [] |
| 55 | |
| 56 | # Tokenise and obtain per-token hidden states |
| 57 | inputs = self.tokenizer( |
| 58 | text, |
| 59 | return_tensors="pt", |
| 60 | return_offsets_mapping=True, |
| 61 | truncation=True, |
| 62 | max_length=self.max_len, |
| 63 | ) |
| 64 | inputs = {k: v.to(self.device) for k, v in inputs.items()} |
| 65 | offsets = inputs.pop("offset_mapping").squeeze(0).cpu().tolist() # (seq_len, 2) |
| 66 | |
| 67 | out = self.model(**inputs) |
| 68 | last_hidden = out.last_hidden_state.squeeze(0) # (seq_len, dim) |
| 69 | last_hidden = last_hidden.cpu() |
| 70 | |
| 71 | # For each chunk span, gather token indices belonging to it |
| 72 | vectors: List[np.ndarray] = [] |
| 73 | for start_char, end_char in chunk_spans: |
| 74 | token_indices = [i for i, (s, e) in enumerate(offsets) if s >= start_char and e <= end_char] |
| 75 | if not token_indices: |
| 76 | # Fallback: if tokenizer lost the span (e.g. due to trimming) just average CLS + SEP |
| 77 | token_indices = [0] |
| 78 | chunk_vec = last_hidden[token_indices].mean(dim=0).numpy().astype("float32") |