(self, model_name: str = "Qwen/Qwen3-Embedding-0.6B", *, max_tokens: int = 8192)
| 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]: |
nothing calls this directly
no outgoing calls
no test coverage detected