| import re |
| from collections import Counter |
| from dataclasses import dataclass |
| from typing import Dict, Iterable, List, Optional, Sequence, Tuple |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| TOKEN_PATTERN = re.compile(r"<image>|[A-Za-z0-9]+(?:[-'][A-Za-z0-9]+)?|[^\sA-Za-z0-9]", re.UNICODE) |
|
|
|
|
| class BasicVQATokenizer: |
| """ |
| Tokenizer tối giản cho prototype VQA classification. |
| |
| Không thay thế tokenizer LLM thật. Mục tiêu là smoke test cross-modal pipeline: |
| question_text -> input_ids/attention_mask -> SimpleTextEncoder. |
| """ |
|
|
| pad_token = "<pad>" |
| unk_token = "<unk>" |
| bos_token = "<bos>" |
| eos_token = "<eos>" |
|
|
| def __init__( |
| self, |
| vocab: Optional[Dict[str, int]] = None, |
| lowercase: bool = True, |
| keep_image_token: bool = True, |
| ): |
| self.lowercase = lowercase |
| self.keep_image_token = keep_image_token |
|
|
| if vocab is None: |
| vocab = { |
| self.pad_token: 0, |
| self.unk_token: 1, |
| self.bos_token: 2, |
| self.eos_token: 3, |
| "<image>": 4, |
| } |
| self.vocab = dict(vocab) |
| self.id_to_token = {idx: token for token, idx in self.vocab.items()} |
|
|
| @property |
| def pad_token_id(self) -> int: |
| return self.vocab[self.pad_token] |
|
|
| @property |
| def unk_token_id(self) -> int: |
| return self.vocab[self.unk_token] |
|
|
| @property |
| def bos_token_id(self) -> int: |
| return self.vocab[self.bos_token] |
|
|
| @property |
| def eos_token_id(self) -> int: |
| return self.vocab[self.eos_token] |
|
|
| def __len__(self) -> int: |
| return len(self.vocab) |
|
|
| def tokenize(self, text: str) -> List[str]: |
| text = text or "" |
| tokens = TOKEN_PATTERN.findall(text) |
| processed = [] |
| for token in tokens: |
| if token == "<image>": |
| if self.keep_image_token: |
| processed.append(token) |
| continue |
| processed.append(token.lower() if self.lowercase else token) |
| return processed |
|
|
| def encode( |
| self, |
| text: str, |
| max_length: int = 64, |
| add_special_tokens: bool = True, |
| ) -> List[int]: |
| tokens = self.tokenize(text) |
| ids = [self.vocab.get(token, self.unk_token_id) for token in tokens] |
|
|
| if add_special_tokens: |
| ids = [self.bos_token_id] + ids + [self.eos_token_id] |
|
|
| return ids[:max_length] |
|
|
| def batch_encode( |
| self, |
| texts: Sequence[str], |
| max_length: int = 64, |
| add_special_tokens: bool = True, |
| ) -> Dict[str, torch.Tensor]: |
| encoded = [self.encode(text, max_length=max_length, add_special_tokens=add_special_tokens) for text in texts] |
| max_len = max((len(ids) for ids in encoded), default=1) |
| max_len = min(max_len, max_length) |
|
|
| input_ids = torch.full((len(texts), max_len), self.pad_token_id, dtype=torch.long) |
| attention_mask = torch.zeros((len(texts), max_len), dtype=torch.long) |
|
|
| for i, ids in enumerate(encoded): |
| ids = ids[:max_len] |
| input_ids[i, : len(ids)] = torch.tensor(ids, dtype=torch.long) |
| attention_mask[i, : len(ids)] = 1 |
|
|
| return { |
| "input_ids": input_ids, |
| "attention_mask": attention_mask, |
| } |
|
|
| @classmethod |
| def build_from_texts( |
| cls, |
| texts: Iterable[str], |
| min_freq: int = 1, |
| max_vocab_size: Optional[int] = None, |
| lowercase: bool = True, |
| keep_image_token: bool = True, |
| ) -> "BasicVQATokenizer": |
| tokenizer = cls(lowercase=lowercase, keep_image_token=keep_image_token) |
| counter = Counter() |
| for text in texts: |
| counter.update(tokenizer.tokenize(text)) |
|
|
| reserved = dict(tokenizer.vocab) |
| vocab = dict(reserved) |
| max_new_tokens = None |
| if max_vocab_size is not None: |
| max_new_tokens = max(max_vocab_size - len(reserved), 0) |
|
|
| added = 0 |
| for token, freq in counter.most_common(): |
| if freq < min_freq: |
| continue |
| if token in vocab: |
| continue |
| if max_new_tokens is not None and added >= max_new_tokens: |
| break |
| vocab[token] = len(vocab) |
| added += 1 |
|
|
| return cls(vocab=vocab, lowercase=lowercase, keep_image_token=keep_image_token) |
|
|
|
|
| @dataclass |
| class SimpleTextEncoderConfig: |
| vocab_size: int |
| d_model: int = 768 |
| max_length: int = 64 |
| pad_token_id: int = 0 |
| num_layers: int = 2 |
| num_heads: int = 8 |
| dim_feedforward: int = 2048 |
| dropout: float = 0.10 |
| use_positional_embedding: bool = True |
|
|
| def __post_init__(self): |
| if self.vocab_size <= 0: |
| raise ValueError("vocab_size phải > 0") |
| if self.d_model <= 0: |
| raise ValueError("d_model phải > 0") |
| if self.max_length <= 0: |
| raise ValueError("max_length phải > 0") |
| if self.d_model % self.num_heads != 0: |
| raise ValueError("d_model phải chia hết cho num_heads") |
|
|
|
|
| class SimpleTextEncoder(nn.Module): |
| """ |
| Text encoder tối giản: Embedding + TransformerEncoder. |
| |
| Input: |
| input_ids: [B, Nt] |
| attention_mask: [B, Nt], 1=valid, 0=pad |
| |
| Output: |
| text_tokens: [B, Nt, D] |
| pooled_text: [B, D] |
| """ |
|
|
| def __init__(self, config: SimpleTextEncoderConfig): |
| super().__init__() |
| self.config = config |
|
|
| self.token_embedding = nn.Embedding( |
| config.vocab_size, |
| config.d_model, |
| padding_idx=config.pad_token_id, |
| ) |
| self.position_embedding = None |
| if config.use_positional_embedding: |
| self.position_embedding = nn.Embedding(config.max_length, config.d_model) |
|
|
| encoder_layer = nn.TransformerEncoderLayer( |
| d_model=config.d_model, |
| nhead=config.num_heads, |
| dim_feedforward=config.dim_feedforward, |
| dropout=config.dropout, |
| activation="gelu", |
| batch_first=True, |
| norm_first=False, |
| ) |
| self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=config.num_layers) |
| self.output_norm = nn.LayerNorm(config.d_model) |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| return_dict: bool = True, |
| ): |
| if input_ids.ndim != 2: |
| raise ValueError(f"input_ids phải có shape [B, Nt], nhận {tuple(input_ids.shape)}") |
|
|
| batch_size, seq_len = input_ids.shape |
| if seq_len > self.config.max_length: |
| input_ids = input_ids[:, : self.config.max_length] |
| seq_len = self.config.max_length |
| if attention_mask is not None: |
| attention_mask = attention_mask[:, : self.config.max_length] |
|
|
| if attention_mask is None: |
| attention_mask = (input_ids != self.config.pad_token_id).long() |
| attention_mask = attention_mask.to(device=input_ids.device) |
|
|
| x = self.token_embedding(input_ids) |
| if self.position_embedding is not None: |
| positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch_size, -1) |
| x = x + self.position_embedding(positions) |
|
|
| key_padding_mask = attention_mask == 0 |
| text_tokens = self.encoder(x, src_key_padding_mask=key_padding_mask) |
| text_tokens = self.output_norm(text_tokens) |
|
|
| pooled_text = self.masked_mean_pool(text_tokens, attention_mask) |
|
|
| if not return_dict: |
| return text_tokens |
|
|
| return { |
| "text_tokens": text_tokens, |
| "pooled_text": pooled_text, |
| "attention_mask": attention_mask, |
| } |
|
|
| @staticmethod |
| def masked_mean_pool(tokens: torch.Tensor, attention_mask: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: |
| mask = attention_mask.to(device=tokens.device, dtype=tokens.dtype).unsqueeze(-1) |
| summed = (tokens * mask).sum(dim=1) |
| denom = mask.sum(dim=1).clamp_min(eps) |
| return summed / denom |
|
|
|
|
| def build_text_encoder_from_questions( |
| questions: Sequence[str], |
| d_model: int = 768, |
| max_length: int = 64, |
| min_freq: int = 1, |
| max_vocab_size: Optional[int] = 30000, |
| num_layers: int = 2, |
| num_heads: int = 8, |
| dropout: float = 0.10, |
| ) -> Tuple[BasicVQATokenizer, SimpleTextEncoder]: |
| tokenizer = BasicVQATokenizer.build_from_texts( |
| questions, |
| min_freq=min_freq, |
| max_vocab_size=max_vocab_size, |
| ) |
| config = SimpleTextEncoderConfig( |
| vocab_size=len(tokenizer), |
| d_model=d_model, |
| max_length=max_length, |
| pad_token_id=tokenizer.pad_token_id, |
| num_layers=num_layers, |
| num_heads=num_heads, |
| dropout=dropout, |
| ) |
| return tokenizer, SimpleTextEncoder(config) |
|
|
