| """ |
| TELEN: Temporal Evolving Legal Embedding Network. |
| |
| Bi-encoder backbone + Legal Concept Graph + HyperNetwork projection. |
| Embedding space adapts dynamically to the legal corpus state. |
| """ |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import AutoModel, AutoTokenizer |
| from pyvi import ViTokenizer |
|
|
| from .config import TELENConfig |
| from .hypernetwork import StateEncoder, HyperNetwork |
| from .concept_graph import build_concept_graph |
|
|
|
|
| def wseg(text): |
| return ViTokenizer.tokenize(text.replace("_", " ")) |
|
|
|
|
| class BiEncoder(nn.Module): |
| """Vietnamese bi-encoder backbone with attention pooling.""" |
|
|
| def __init__(self, model_name="bkai-foundation-models/vietnamese-bi-encoder"): |
| super().__init__() |
| self.model = AutoModel.from_pretrained(model_name) |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) |
| self.dim = self.model.config.hidden_size |
| self.attn_query = nn.Parameter(torch.randn(self.dim)) |
| self.scale = self.dim ** 0.5 |
|
|
| def forward(self, texts, max_len=480): |
| segmented = [wseg(t) for t in texts] |
| enc = self.tokenizer(segmented, padding=True, truncation=True, |
| max_length=max_len, return_tensors="pt") |
| input_ids = enc["input_ids"].to(self.attn_query.device) |
| mask = enc["attention_mask"].to(self.attn_query.device) |
| hidden = self.model(input_ids=input_ids, attention_mask=mask).last_hidden_state |
| scores = torch.einsum("bsd,d->bs", hidden, self.attn_query) / self.scale |
| scores = scores.masked_fill(mask == 0, float("-1e9")) |
| weights = F.softmax(scores, dim=1) |
| return torch.einsum("bsd,bs->bd", hidden, weights) |
|
|
|
|
| class TELEN(nn.Module): |
| """Temporal Evolving Legal Embedding Network.""" |
|
|
| def __init__(self, config: TELENConfig): |
| super().__init__() |
| self.config = config |
| d = config.hidden_dim |
|
|
| |
| self.encoder = BiEncoder() |
| for p in self.encoder.parameters(): |
| p.requires_grad = False |
|
|
| |
| self.base_projection = nn.Sequential(nn.Linear(d, d), nn.Tanh()) |
| self.proj_norm = nn.LayerNorm(d) |
| self.attn_query = nn.Parameter(torch.randn(d)) |
|
|
| |
| self.concept_graph = None |
| self.law_id_to_idx = None |
|
|
| |
| self.state_encoder = StateEncoder(d) |
| self.hypernetwork = HyperNetwork(config) |
|
|
| def _pool(self, hidden, mask): |
| """Attention-weighted pooling (for pre-tokenized inputs).""" |
| scores = torch.einsum("bsd,d->bs", hidden, self.attn_query) / (self.config.hidden_dim ** 0.5) |
| scores = scores.masked_fill(mask == 0, float("-1e9")) |
| weights = F.softmax(scores, dim=1) |
| return torch.einsum("bsd,bs->bd", hidden, weights) |
|
|
| def encode_text(self, texts): |
| return self.encoder(texts, max_len=self.config.max_seq_length) |
|
|
| def get_state_vector(self): |
| if self.concept_graph is None or self.concept_graph.num_nodes == 0: |
| return torch.zeros(self.config.hidden_dim, device=self.attn_query.device) |
| refined = self.concept_graph.forward() |
| return self.state_encoder(refined) |
|
|
| def adapt_embedding(self, raw, state_vec): |
| base = self.base_projection(raw) |
| hn = self.hypernetwork(state_vec) |
| shift = raw @ hn["shift_matrix"].T + hn["bias"] |
| mean = F.normalize(self.proj_norm(base + shift), p=2, dim=1) |
| result = {"mean": mean, "log_variance": hn.get("log_variance")} |
| if self.config.hypernetwork.output_variance: |
| noise = 0.1 * hn["log_variance"].exp().clamp(min=0.001, max=0.25).sqrt().clamp(max=0.5) |
| result["sample"] = F.normalize(mean + torch.randn_like(mean) * noise, p=2, dim=1) |
| else: |
| result["sample"] = mean |
| return result |
|
|
| def forward(self, texts, use_stochastic=False): |
| raw = self.encode_text(texts) |
| state = self.get_state_vector() |
| adapted = self.adapt_embedding(raw, state) |
| return { |
| "embeddings": adapted["sample"] if use_stochastic else adapted["mean"], |
| "mean": adapted["mean"], |
| "log_variance": adapted.get("log_variance"), |
| "state_vector": state, |
| } |
|
|
| def build_graph(self, df): |
| self.concept_graph, self.law_id_to_idx = build_concept_graph( |
| df, lambda t: self.encode_text([t])[0].detach(), self.config, |
| ) |
| self.concept_graph = self.concept_graph.to(self.attn_query.device) |
|
|
| def add_law(self, law_id, articles): |
| if self.concept_graph is None: return |
| if articles: |
| emb = self.encode_text(articles[:5]).mean(dim=0) |
| new_idx = self.concept_graph.num_nodes |
| self.concept_graph.add_nodes([law_id], emb.unsqueeze(0)) |
| existing = self.concept_graph.node_embeddings[:-1] |
| if len(existing) > 0: |
| sim = F.cosine_similarity(emb.unsqueeze(0), existing) |
| _, top = sim.topk(k=min(10, len(existing))) |
| self.concept_graph.add_edges("semantic", |
| [(new_idx, i.item(), sim[i].item()) for i in top]) |
|
|
|
|
| def create_model(config: TELENConfig) -> TELEN: |
| return TELEN(config) |
|
|
