Spaces:
Running
Running
| """ | |
| CoMER model: DenseNet Encoder + Transformer Decoder with ARM. | |
| Based on: https://github.com/Green-Wood/CoMER | |
| Stripped from ICAL: single CE loss, no SCCM/FusionModule. | |
| """ | |
| from typing import List | |
| import torch | |
| import torch.nn as nn | |
| from torch import FloatTensor, LongTensor | |
| from .decoder import Decoder | |
| from .encoder import Encoder | |
| class CoMER(nn.Module): | |
| def __init__( | |
| self, | |
| d_model: int, | |
| growth_rate: int, | |
| num_layers: int, | |
| nhead: int, | |
| num_decoder_layers: int, | |
| dim_feedforward: int, | |
| dropout: float, | |
| dc: int, | |
| cross_coverage: bool, | |
| self_coverage: bool, | |
| vocab_size: int, | |
| pad_idx: int = 0, | |
| ): | |
| super().__init__() | |
| self.encoder = Encoder( | |
| d_model=d_model, growth_rate=growth_rate, num_layers=num_layers | |
| ) | |
| self.decoder = Decoder( | |
| d_model=d_model, | |
| nhead=nhead, | |
| num_decoder_layers=num_decoder_layers, | |
| dim_feedforward=dim_feedforward, | |
| dropout=dropout, | |
| dc=dc, | |
| cross_coverage=cross_coverage, | |
| self_coverage=self_coverage, | |
| vocab_size=vocab_size, | |
| pad_idx=pad_idx, | |
| ) | |
| def forward( | |
| self, img: FloatTensor, img_mask: LongTensor, tgt: LongTensor | |
| ) -> FloatTensor: | |
| """ | |
| Parameters | |
| ---------- | |
| img : FloatTensor [b, 1, h, w] | |
| img_mask: LongTensor [b, h, w] | |
| tgt : LongTensor [b, l] | |
| Returns | |
| ------- | |
| FloatTensor [b, l, vocab_size] | |
| """ | |
| feature, mask = self.encoder(img, img_mask) # [b, h, w, d] | |
| out = self.decoder(feature, mask, tgt) | |
| return out | |
| def encode(self, img: FloatTensor, img_mask: LongTensor): | |
| """Encode image only (for beam search).""" | |
| return self.encoder(img, img_mask) | |
| def greedy_decode( | |
| self, | |
| img: FloatTensor, | |
| img_mask: LongTensor, | |
| sos_idx: int, | |
| eos_idx: int, | |
| max_len: int = 200, | |
| ) -> List[List[int]]: | |
| """Greedy decode for validation (l2r only, no bidirectional). | |
| Parameters | |
| ---------- | |
| img : FloatTensor [b, 1, h, w] | |
| img_mask : LongTensor [b, h, w] | |
| Returns | |
| ------- | |
| List[List[int]]: decoded token indices (no SOS/EOS) | |
| """ | |
| self.eval() | |
| b = img.size(0) | |
| feature, mask = self.encoder(img, img_mask) | |
| input_ids = torch.full( | |
| (b, 1), fill_value=sos_idx, dtype=torch.long, device=img.device | |
| ) | |
| finished = torch.zeros(b, dtype=torch.bool, device=img.device) | |
| for _ in range(max_len): | |
| out = self.decoder(feature, mask, input_ids) # [b, l, vocab] | |
| next_token = out[:, -1, :].argmax(dim=-1) # [b] | |
| finished = finished | (next_token == eos_idx) | |
| next_token[finished] = eos_idx | |
| input_ids = torch.cat( | |
| [input_ids, next_token.unsqueeze(1)], dim=1 | |
| ) | |
| if finished.all(): | |
| break | |
| results = [] | |
| for i in range(b): | |
| seq = input_ids[i, 1:].tolist() | |
| if eos_idx in seq: | |
| seq = seq[:seq.index(eos_idx)] | |
| results.append(seq) | |
| return results | |
| def beam_search_decode( | |
| self, | |
| img: FloatTensor, | |
| img_mask: LongTensor, | |
| sos_idx: int, | |
| eos_idx: int, | |
| pad_idx: int = 0, | |
| beam_size: int = 10, | |
| max_len: int = 200, | |
| alpha: float = 1.0, | |
| ) -> List[List[int]]: | |
| """Beam search decode (l2r only, clean implementation). | |
| Processes one sample at a time to avoid complexity of batched beam search. | |
| """ | |
| import torch.nn.functional as F | |
| self.eval() | |
| b = img.size(0) | |
| device = img.device | |
| feature, mask = self.encoder(img, img_mask) # [b, h, w, d] | |
| results = [] | |
| for i in range(b): | |
| feat_i = feature[i:i+1] # [1, h, w, d] | |
| mask_i = mask[i:i+1] # [1, h, w] | |
| # Expand for beams: [beam, h, w, d] | |
| feat_beam = feat_i.expand(beam_size, -1, -1, -1) | |
| mask_beam = mask_i.expand(beam_size, -1, -1) | |
| # Each beam starts with SOS | |
| sequences = torch.full((beam_size, 1), sos_idx, dtype=torch.long, device=device) | |
| scores = torch.zeros(beam_size, device=device) | |
| scores[1:] = -1e9 # only first beam active initially | |
| # Store completed hypotheses: (score, sequence) | |
| completed = [] | |
| for step in range(max_len): | |
| out = self.decoder(feat_beam, mask_beam, sequences) # [beam, l, vocab] | |
| logits = out[:, -1, :] # [beam, vocab] | |
| log_probs = F.log_softmax(logits, dim=-1) | |
| # Candidate scores: [beam, vocab] | |
| candidate_scores = scores.unsqueeze(-1) + log_probs | |
| vocab_size = log_probs.size(-1) | |
| if step == 0: | |
| # Only first beam is active | |
| candidate_scores = candidate_scores[0:1].reshape(-1) | |
| else: | |
| candidate_scores = candidate_scores.reshape(-1) | |
| # Top-k candidates | |
| topk_scores, topk_ids = candidate_scores.topk(beam_size, dim=-1) | |
| beam_ids = topk_ids // vocab_size | |
| token_ids = topk_ids % vocab_size | |
| # Build new sequences | |
| new_sequences = torch.cat([ | |
| sequences[beam_ids], token_ids.unsqueeze(-1) | |
| ], dim=-1) | |
| new_scores = topk_scores | |
| # Check for completed beams (EOS generated) | |
| active_mask = token_ids != eos_idx | |
| for j in range(beam_size): | |
| if not active_mask[j]: | |
| seq = new_sequences[j, 1:].tolist() # remove SOS | |
| if eos_idx in seq: | |
| seq = seq[:seq.index(eos_idx)] | |
| # Length-normalized score | |
| length_penalty = ((5.0 + len(seq)) / 6.0) ** alpha | |
| normalized_score = new_scores[j].item() / length_penalty | |
| completed.append((normalized_score, seq)) | |
| # Keep only active beams | |
| active_indices = active_mask.nonzero(as_tuple=True)[0] | |
| if len(active_indices) == 0: | |
| break | |
| if len(completed) >= beam_size: | |
| break | |
| # Pad back to beam_size if needed | |
| if len(active_indices) < beam_size: | |
| # Fill remaining slots with top active beams | |
| pad_count = beam_size - len(active_indices) | |
| pad_indices = active_indices[:pad_count].repeat( | |
| (pad_count + len(active_indices) - 1) // len(active_indices) + 1 | |
| )[:pad_count] | |
| all_indices = torch.cat([active_indices, pad_indices]) | |
| new_scores_padded = new_scores[all_indices] | |
| new_scores_padded[len(active_indices):] = -1e9 | |
| sequences = new_sequences[all_indices] | |
| scores = new_scores_padded | |
| else: | |
| sequences = new_sequences[active_indices[:beam_size]] | |
| scores = new_scores[active_indices[:beam_size]] | |
| # If no completed, use best active beam | |
| if not completed: | |
| seq = sequences[0, 1:].tolist() | |
| if eos_idx in seq: | |
| seq = seq[:seq.index(eos_idx)] | |
| completed.append((scores[0].item(), seq)) | |
| # Pick best hypothesis | |
| completed.sort(key=lambda x: x[0], reverse=True) | |
| results.append(completed[0][1]) | |
| return results | |
| def build_model(config, vocab_size: int) -> CoMER: | |
| """Build CoMER model from config.""" | |
| model = CoMER( | |
| d_model=config.model.d_model, | |
| growth_rate=config.model.growth_rate, | |
| num_layers=config.model.num_layers, | |
| nhead=config.model.nhead, | |
| num_decoder_layers=config.model.num_decoder_layers, | |
| dim_feedforward=config.model.dim_feedforward, | |
| dropout=config.model.dropout, | |
| dc=config.model.dc, | |
| cross_coverage=config.model.cross_coverage, | |
| self_coverage=config.model.self_coverage, | |
| vocab_size=vocab_size, | |
| pad_idx=0, | |
| ) | |
| return model | |