File size: 4,718 Bytes
bff06c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """GreenLeaf Law Embed — model implementation.
Bidirectional transformer encoder for legal-domain text embedding.
Built on the Qwen3 architecture with causal masking removed,
enabling full-sequence attention in both directions.
This is critical for legal text where relevant context (holdings,
citations, defined terms) can appear anywhere in a document.
"""
import inspect
from typing import Callable
import torch
from transformers import Qwen3Model
from transformers.cache_utils import Cache
from transformers.masking_utils import create_causal_mask
from transformers.modeling_outputs import BaseModelOutputWithPooling
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs
from .configuration import GreenLeafEmbedConfig
# ---------------------------------------------------------------------------
# Compatibility shim: the `create_causal_mask` API has changed across
# transformers releases. We detect the correct parameter names at import
# time so this code works across versions 5.1 through 5.15+.
#
# <= 5.1 : kwarg is `input_embeds`, `cache_position` required
# 5.2-5.5 : renamed to `inputs_embeds`, `cache_position` still required
# 5.6-5.8 : `cache_position` has a default (backward compat)
# >= 5.9 : `cache_position` removed entirely
# ---------------------------------------------------------------------------
_mask_fn_params = inspect.signature(create_causal_mask).parameters
_embeds_param = "inputs_embeds" if "inputs_embeds" in _mask_fn_params else "input_embeds"
_has_cache_position = "cache_position" in _mask_fn_params
def _build_bidirectional_mask_fn(attn_mask: torch.Tensor | None) -> Callable:
"""Return a mask function that allows every token to attend to every
other token, subject only to the padding mask.
Standard causal masking restricts token i to attend only to tokens
j <= i. For embedding models we want the opposite — full visibility —
so that the representation of each token is informed by the entire
input sequence.
"""
def _mask(batch: int, head: int, q_pos: int, kv_pos: int) -> bool:
if attn_mask is None:
return torch.ones((), dtype=torch.bool)
return attn_mask[batch, kv_pos].to(torch.bool)
return _mask
class GreenLeafEmbedModel(Qwen3Model):
"""Bidirectional Qwen3 encoder for text embedding.
Overrides the causal self-attention in Qwen3 with bidirectional
attention so that every token representation captures full-sequence
context. This is essential for retrieval tasks where the meaning of
a passage depends on information that may appear before or after
any given token.
"""
_supports_flash_attn = True
_supports_sdpa = True
config_class = GreenLeafEmbedConfig
def __init__(self, config):
super().__init__(config)
self.post_init()
def post_init(self):
super().post_init()
# Disable causal attention in every transformer layer.
# This works with both flash_attention_2 and sdpa backends.
for layer in self.layers:
layer.self_attn.is_causal = False
def forward(
self,
input_ids: torch.LongTensor | None = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: Cache | None = None,
inputs_embeds: torch.FloatTensor | None = None,
use_cache: bool | None = None,
cache_position: torch.LongTensor | None = None,
**kwargs: Unpack[TransformersKwargs],
) -> BaseModelOutputWithPooling:
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
input_ids = None
mask_args = {
"config": self.config,
_embeds_param: inputs_embeds,
"attention_mask": attention_mask,
"past_key_values": None,
"position_ids": position_ids,
"or_mask_function": _build_bidirectional_mask_fn(attention_mask),
}
if _has_cache_position:
mask_args["cache_position"] = torch.arange(
inputs_embeds.shape[1],
device=inputs_embeds.device,
dtype=torch.long,
)
attention_mask = {"full_attention": create_causal_mask(**mask_args)}
outputs = super().forward(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
cache_position=cache_position,
**kwargs,
)
return outputs
|