| """ |
| objective.py — Objective Layer untuk AKSARA. |
| |
| Komponen: |
| embedding_relation_loss : cek apakah verb dan objek "nyambung" secara semantik |
| berbasis semantic_slots dari BSU — BUKAN embedding baru |
| cooccurrence_loss : cek apakah kombinasi token masuk akal secara statistik |
| berbasis PMI matrix yang diprecompute dari corpus |
| build_cooccurrence_matrix: precompute PMI(a, b) dari corpus, domain-agnostic |
| negative_sample : buat kalimat negatif via token shuffle |
| |
| Prinsip implementasi: |
| - Tidak ada hardcode domain atau aturan semantik manual |
| - Semua berbasis representasi model sendiri (BSU semantic_slots) |
| - Co-occurrence berbasis data corpus, bukan library NLP eksternal |
| - Tidak ada layer baru — hanya operasi di atas output BSU yang sudah ada |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| import random |
| from collections import defaultdict |
| from typing import Dict, List, Optional, Set, Tuple |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def build_cooccurrence_matrix( |
| corpus: List[str], |
| root_vocab: Dict[str, int], |
| window: int = 4, |
| min_count: int = 2, |
| ) -> Dict[Tuple[int, int], float]: |
| token_count: Dict[int, int] = defaultdict(int) |
| pair_count: Dict[Tuple[int, int], int] = defaultdict(int) |
| total_tokens = 0 |
|
|
| for text in corpus: |
| tokens = text.lower().split() |
| ids = [root_vocab[t] for t in tokens if t in root_vocab] |
| for i, a in enumerate(ids): |
| token_count[a] += 1 |
| total_tokens += 1 |
| for j in range(i + 1, min(i + 1 + window, len(ids))): |
| b = ids[j] |
| if a != b: |
| pair_count[(min(a, b), max(a, b))] += 1 |
|
|
| if total_tokens == 0: |
| return {} |
|
|
| co_matrix: Dict[Tuple[int, int], float] = {} |
| for (a, b), co_c in pair_count.items(): |
| if co_c < min_count: |
| continue |
| p_ab = co_c / total_tokens |
| p_a = token_count[a] / total_tokens |
| p_b = token_count[b] / total_tokens |
| pmi = math.log(p_ab / (p_a * p_b + 1e-9) + 1e-9) |
| if pmi > 0: |
| co_matrix[(a, b)] = pmi |
|
|
| return co_matrix |
|
|
|
|
| def embedding_relation_loss( |
| semantic_slots: torch.Tensor, |
| morpheme_ids: torch.Tensor, |
| verb_token_ids: Set[int], |
| min_sim_target: float = 0.3, |
| ) -> torch.Tensor: |
| device = semantic_slots.device |
| B, L, d = semantic_slots.shape |
|
|
| if L < 2: |
| return torch.tensor(0.0, device=device) |
|
|
| total_loss = torch.tensor(0.0, device=device) |
| n_valid = 0 |
|
|
| for b in range(B): |
| ids = morpheme_ids[b] |
| sem = semantic_slots[b] |
|
|
| verb_pos = -1 |
| for pos in range(L): |
| if ids[pos].item() in verb_token_ids: |
| verb_pos = pos |
| break |
|
|
| if verb_pos == -1 or verb_pos >= L - 1: |
| continue |
|
|
| v_sem = sem[verb_pos] |
| after_verb = sem[verb_pos + 1:] |
| if after_verb.size(0) == 0: |
| continue |
|
|
| v_norm = F.normalize(v_sem.unsqueeze(0), dim=-1) |
| o_norm = F.normalize(after_verb, dim=-1) |
| sims = (v_norm * o_norm).sum(dim=-1) |
|
|
| avg_sim = sims.mean() |
| deficit = (min_sim_target - avg_sim).clamp(min=0.0) |
| total_loss = total_loss + deficit |
| n_valid += 1 |
|
|
| if n_valid == 0: |
| return torch.tensor(0.0, device=device) |
|
|
| return total_loss / n_valid |
|
|
|
|
| def cooccurrence_loss( |
| morpheme_ids: torch.Tensor, |
| co_matrix: Dict[Tuple[int, int], float], |
| window: int = 3, |
| ) -> torch.Tensor: |
| device = morpheme_ids.device |
| B, L = morpheme_ids.shape |
|
|
| total_pmi = 0.0 |
| n_pairs = 0 |
|
|
| for b in range(B): |
| ids = morpheme_ids[b].tolist() |
| for i in range(L): |
| a = ids[i] |
| if a == 0: |
| continue |
| for j in range(i + 1, min(i + 1 + window, L)): |
| bb = ids[j] |
| if bb == 0: |
| continue |
| key = (min(a, bb), max(a, bb)) |
| if key in co_matrix: |
| total_pmi += co_matrix[key] |
| n_pairs += 1 |
|
|
| if n_pairs == 0: |
| return torch.tensor(0.0, device=device) |
|
|
| avg_pmi = total_pmi / n_pairs |
| return torch.tensor(-avg_pmi, dtype=torch.float32, device=device) |
|
|
|
|
| def make_negative_batch( |
| morpheme_ids: torch.Tensor, |
| n_neg: int = 1, |
| ) -> torch.Tensor: |
| B, L = morpheme_ids.shape |
| device = morpheme_ids.device |
| negs = [] |
|
|
| for _ in range(n_neg): |
| neg = morpheme_ids.clone() |
| for b in range(B): |
| non_pad = (neg[b] != 0).nonzero(as_tuple=True)[0] |
| if len(non_pad) > 2: |
| perm = non_pad[torch.randperm(len(non_pad), device=device)] |
| neg[b, non_pad] = neg[b, perm] |
| negs.append(neg) |
|
|
| return torch.cat(negs, dim=0) |
|
|
|
|
| class CompositeLoss: |
| def __init__( |
| self, |
| lambda_rel: float = 0.3, |
| lambda_co: float = 0.1, |
| verb_ids: Optional[Set[int]] = None, |
| co_matrix: Optional[Dict[Tuple[int, int], float]] = None, |
| ): |
| self.lambda_rel = lambda_rel |
| self.lambda_co = lambda_co |
| self.verb_ids = verb_ids or set() |
| self.co_matrix = co_matrix or {} |
|
|
| def __call__( |
| self, |
| base_loss: torch.Tensor, |
| model_output: Dict, |
| morpheme_ids: torch.Tensor, |
| ) -> Tuple[torch.Tensor, Dict[str, float]]: |
| breakdown: Dict[str, float] = {"base": base_loss.item()} |
| total = base_loss |
|
|
| semantic_slots = model_output.get("semantic_slots") |
| if semantic_slots is None: |
| semantic_slots = model_output.get("bsu_original") |
|
|
| if semantic_slots is not None and self.verb_ids: |
| l_rel = embedding_relation_loss( |
| semantic_slots, morpheme_ids, self.verb_ids, |
| ) |
| total = total + self.lambda_rel * l_rel |
| breakdown["rel"] = l_rel.item() |
| else: |
| breakdown["rel"] = 0.0 |
|
|
| if self.co_matrix: |
| l_co = cooccurrence_loss(morpheme_ids, self.co_matrix) |
| total = total + self.lambda_co * l_co |
| breakdown["co"] = l_co.item() |
| else: |
| breakdown["co"] = 0.0 |
|
|
| breakdown["total"] = total.item() |
| return total, breakdown |
