Replace broken LRSCLIPTextEncoder with official DGTRS architecture
Root cause of NaN: our open_clip wrapper had 3 bugs: 1. Positional embeddings summed for all positions instead of masked (official: mask1 for pos 0-19, mask2 for pos 20-247) 2. open_clip uses batch-first transformer, DGTRS uses sequence-first (LND format with nn.MultiheadAttention) 3. open_clip tokenizer truncates to 77 tokens, DGTRS needs 248 Fix: copied official DGTRS text encoder architecture from github.com/MitsuiChen14/DGTRS (Apache-2.0): - src/models/dgtrs/model.py: DGTRSTextEncoder, build_model, load_dgtrs_text_encoder, tokenize_dgtrs - src/models/dgtrs/simple_tokenizer.py: BPE tokenizer (248 tokens) - src/models/dgtrs/bpe_simple_vocab_16e6.txt.gz: vocabulary Removed: LRSCLIPTextEncoder class, open_clip dependency for text encoding Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
@@ -17,7 +17,6 @@ import warnings
|
||||
from pathlib import Path
|
||||
|
||||
import coloredlogs
|
||||
import open_clip
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
@@ -26,6 +25,7 @@ LOGGER = logging.getLogger("caption_test.model")
|
||||
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||||
from safetensors.torch import load_file as load_safetensors
|
||||
|
||||
from src.models.dgtrs.model import DGTRSTextEncoder, load_dgtrs_text_encoder, tokenize_dgtrs
|
||||
from src.models.dual_encoder import GatedFusion, ProjectionHead
|
||||
|
||||
|
||||
@@ -197,116 +197,8 @@ class DINOv3ViT(nn.Module):
|
||||
return model
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LRSCLIP (DGTRS-CLIP) text encoder — open_clip with KPS positional embedding
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class LRSCLIPTextEncoder(nn.Module):
|
||||
"""DGTRS-CLIP text encoder with KPS positional embedding (248 tokens).
|
||||
|
||||
Wraps open_clip ViT-L-14 text tower. Handles the extra
|
||||
positional_embedding_res (KPS residual) not present in vanilla open_clip.
|
||||
|
||||
Checkpoint keys (text-only):
|
||||
token_embedding.weight: [49408, 768]
|
||||
positional_embedding: [248, 768]
|
||||
positional_embedding_res: [248, 768]
|
||||
transformer.resblocks.{0-11}.{attn,mlp,ln}.*
|
||||
ln_final.{weight,bias}
|
||||
text_projection: [768, 768]
|
||||
"""
|
||||
|
||||
def __init__(self, context_length: int = 248, embed_dim: int = 768) -> None:
|
||||
super().__init__()
|
||||
self.context_length = context_length
|
||||
self.embed_dim = embed_dim
|
||||
|
||||
# Build the open_clip model to get correct architecture, then
|
||||
# replace positional embedding and add KPS residual.
|
||||
# Suppress "no pretrained weights" warning — we load DGTRS-CLIP weights after.
|
||||
_root = logging.getLogger()
|
||||
_prev_level = _root.level
|
||||
_root.setLevel(logging.ERROR)
|
||||
clip_model = open_clip.create_model("ViT-L-14", pretrained=None)
|
||||
_root.setLevel(_prev_level)
|
||||
|
||||
# Extract text components.
|
||||
self.token_embedding = clip_model.token_embedding
|
||||
self.transformer = clip_model.transformer
|
||||
self.ln_final = clip_model.ln_final
|
||||
self.text_projection = clip_model.text_projection
|
||||
self.attn_mask = None # rebuilt in forward
|
||||
|
||||
# Replace positional embedding with 248-length version.
|
||||
self.positional_embedding = nn.Parameter(
|
||||
torch.zeros(context_length, embed_dim),
|
||||
)
|
||||
self.positional_embedding_res = nn.Parameter(
|
||||
torch.zeros(context_length, embed_dim),
|
||||
)
|
||||
|
||||
def _build_attn_mask(self, context_length: int, device: torch.device) -> torch.Tensor:
|
||||
mask = torch.empty(context_length, context_length, device=device)
|
||||
mask.fill_(float("-inf"))
|
||||
mask.triu_(1)
|
||||
return mask
|
||||
|
||||
def forward(self, text: torch.Tensor) -> torch.Tensor:
|
||||
"""Encode tokenized text.
|
||||
|
||||
Args:
|
||||
text: Token IDs [B, T] (T <= 248).
|
||||
|
||||
Returns:
|
||||
Text embeddings [B, 768], L2-normalized.
|
||||
"""
|
||||
T = text.shape[1]
|
||||
x = self.token_embedding(text) # [B, T, 768]
|
||||
pos = self.positional_embedding[:T] + self.positional_embedding_res[:T]
|
||||
x = x + pos
|
||||
|
||||
if self.attn_mask is None or self.attn_mask.shape[0] != T:
|
||||
self.attn_mask = self._build_attn_mask(T, x.device)
|
||||
attn_mask = self.attn_mask[:T, :T]
|
||||
|
||||
# open_clip transformer expects batch-first [B, T, D].
|
||||
x = self.transformer(x, attn_mask=attn_mask)
|
||||
x = self.ln_final(x) # [B, T, D]
|
||||
|
||||
# Take features at EOS token position.
|
||||
eos_idx = text.argmax(dim=-1)
|
||||
x = x[torch.arange(x.shape[0], device=x.device), eos_idx]
|
||||
|
||||
# Project.
|
||||
if isinstance(self.text_projection, nn.Parameter):
|
||||
x = x @ self.text_projection
|
||||
else:
|
||||
x = self.text_projection(x)
|
||||
|
||||
return F.normalize(x, dim=-1)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, path: str | Path) -> LRSCLIPTextEncoder:
|
||||
"""Load DGTRS-CLIP checkpoint, extracting only text encoder weights."""
|
||||
model = cls()
|
||||
path = Path(path)
|
||||
LOGGER.info("📝 Loading LRSCLIP text encoder from %s", path.name)
|
||||
full_state = torch.load(str(path), map_location="cpu", weights_only=False)
|
||||
if "state_dict" in full_state:
|
||||
full_state = full_state["state_dict"]
|
||||
|
||||
# Filter out visual.* keys — keep only text encoder.
|
||||
text_state = {
|
||||
k: v for k, v in full_state.items()
|
||||
if not k.startswith("visual.")
|
||||
}
|
||||
|
||||
# Handle text_projection shape: checkpoint may be [768, 768] Parameter
|
||||
# while open_clip stores it as nn.Parameter directly.
|
||||
model.load_state_dict(text_state, strict=False)
|
||||
n_params = sum(p.numel() for p in model.parameters())
|
||||
LOGGER.info("📝 LRSCLIP loaded: %s params, context=%d tokens", f"{n_params:,}", model.context_length)
|
||||
return model
|
||||
# LRSCLIPTextEncoder removed — replaced by official DGTRS architecture
|
||||
# in src/models/dgtrs/model.py (DGTRSTextEncoder)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -396,15 +288,13 @@ class AsymmetricEncoder(nn.Module):
|
||||
self._freeze(self.drone_encoder)
|
||||
self._freeze(self.sat_encoder)
|
||||
|
||||
# Text encoder (partial unfreeze).
|
||||
# Text encoder — official DGTRS architecture (partial unfreeze).
|
||||
if not baseline_mode:
|
||||
self.text_encoder = LRSCLIPTextEncoder.from_pretrained(lrsclip_path)
|
||||
self.text_encoder = load_dgtrs_text_encoder(lrsclip_path)
|
||||
self._freeze(self.text_encoder)
|
||||
self._unfreeze_text_last_block()
|
||||
self.tokenizer = open_clip.get_tokenizer("ViT-L-14")
|
||||
else:
|
||||
self.text_encoder = None
|
||||
self.tokenizer = None
|
||||
|
||||
# Projection heads.
|
||||
self.proj_drone = ProjectionHead(
|
||||
@@ -473,12 +363,8 @@ class AsymmetricEncoder(nn.Module):
|
||||
return F.normalize(fused, dim=-1)
|
||||
|
||||
def _encode_single_text(self, texts: list[str]) -> torch.Tensor:
|
||||
"""Tokenize and encode a list of strings."""
|
||||
tokens = self.tokenizer(list(texts)).to(self.device)
|
||||
# Pad/truncate to context_length.
|
||||
T = tokens.shape[1]
|
||||
if T > self.text_encoder.context_length:
|
||||
tokens = tokens[:, :self.text_encoder.context_length]
|
||||
"""Tokenize and encode a list of strings using DGTRS tokenizer."""
|
||||
tokens = tokenize_dgtrs(list(texts)).to(self.device)
|
||||
return self.text_encoder(tokens)
|
||||
|
||||
def forward(
|
||||
|
||||
7
src/models/dgtrs/__init__.py
Normal file
7
src/models/dgtrs/__init__.py
Normal file
@@ -0,0 +1,7 @@
|
||||
"""DGTRS-CLIP (LRSCLIP) text encoder.
|
||||
|
||||
Adapted from https://github.com/MitsuiChen14/DGTRS (Apache-2.0).
|
||||
Only the text encoder components are used.
|
||||
"""
|
||||
|
||||
from src.models.dgtrs.model import build_model, load_dgtrs_text_encoder, tokenize_dgtrs
|
||||
BIN
src/models/dgtrs/bpe_simple_vocab_16e6.txt.gz
Normal file
BIN
src/models/dgtrs/bpe_simple_vocab_16e6.txt.gz
Normal file
Binary file not shown.
277
src/models/dgtrs/model.py
Normal file
277
src/models/dgtrs/model.py
Normal file
@@ -0,0 +1,277 @@
|
||||
"""DGTRS-CLIP model — text encoder components.
|
||||
|
||||
Adapted from https://github.com/MitsuiChen14/DGTRS (Apache-2.0).
|
||||
Only text-encoder classes are kept; vision encoder removed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections import OrderedDict
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import coloredlogs
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from src.models.dgtrs.simple_tokenizer import SimpleTokenizer as _Tokenizer
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.dgtrs")
|
||||
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||||
|
||||
_tokenizer = _Tokenizer()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transformer blocks (DGTRS original — sequence-first LND format)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class LayerNorm(nn.LayerNorm):
|
||||
"""LayerNorm that handles fp16 by casting to fp32 internally."""
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
orig_type = x.dtype
|
||||
ret = super().forward(x.type(torch.float32))
|
||||
return ret.type(orig_type)
|
||||
|
||||
|
||||
class QuickGELU(nn.Module):
|
||||
def forward(self, x: torch.Tensor):
|
||||
return x * torch.sigmoid(1.702 * x)
|
||||
|
||||
|
||||
class ResidualAttentionBlock(nn.Module):
|
||||
def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
|
||||
super().__init__()
|
||||
self.attn = nn.MultiheadAttention(d_model, n_head)
|
||||
self.ln_1 = LayerNorm(d_model)
|
||||
self.mlp = nn.Sequential(OrderedDict([
|
||||
("c_fc", nn.Linear(d_model, d_model * 4)),
|
||||
("gelu", QuickGELU()),
|
||||
("c_proj", nn.Linear(d_model * 4, d_model)),
|
||||
]))
|
||||
self.ln_2 = LayerNorm(d_model)
|
||||
self.attn_mask = attn_mask
|
||||
|
||||
def attention(self, x: torch.Tensor):
|
||||
self.attn_mask = (
|
||||
self.attn_mask.to(dtype=x.dtype, device=x.device)
|
||||
if self.attn_mask is not None else None
|
||||
)
|
||||
return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
x = x + self.attention(self.ln_1(x))
|
||||
x = x + self.mlp(self.ln_2(x))
|
||||
return x
|
||||
|
||||
|
||||
class Transformer(nn.Module):
|
||||
def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
|
||||
super().__init__()
|
||||
self.width = width
|
||||
self.layers = layers
|
||||
self.resblocks = nn.Sequential(
|
||||
*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)]
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
return self.resblocks(x)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DGTRS-CLIP text encoder
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class DGTRSTextEncoder(nn.Module):
|
||||
"""DGTRS-CLIP text encoder with KPS dual positional embeddings.
|
||||
|
||||
Context length: 248 tokens.
|
||||
Uses mask1 (positions 0-19) for positional_embedding and
|
||||
mask2 (positions 20-247) for positional_embedding_res.
|
||||
|
||||
This is the official architecture from
|
||||
https://github.com/MitsuiChen14/DGTRS/blob/main/model/model_longclip.py
|
||||
"""
|
||||
|
||||
CONTEXT_LENGTH = 248
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size: int = 49408,
|
||||
transformer_width: int = 768,
|
||||
transformer_heads: int = 12,
|
||||
transformer_layers: int = 12,
|
||||
embed_dim: int = 768,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.context_length = self.CONTEXT_LENGTH
|
||||
self.embed_dim = embed_dim
|
||||
|
||||
self.token_embedding = nn.Embedding(vocab_size, transformer_width)
|
||||
self.positional_embedding = nn.Parameter(
|
||||
torch.empty(self.CONTEXT_LENGTH, transformer_width),
|
||||
)
|
||||
self.positional_embedding_res = nn.Parameter(
|
||||
torch.empty(self.CONTEXT_LENGTH, transformer_width),
|
||||
)
|
||||
|
||||
self.transformer = Transformer(
|
||||
width=transformer_width,
|
||||
layers=transformer_layers,
|
||||
heads=transformer_heads,
|
||||
attn_mask=self._build_attention_mask(),
|
||||
)
|
||||
|
||||
self.ln_final = LayerNorm(transformer_width)
|
||||
self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
|
||||
|
||||
# KPS masks: positions 0-19 use positional_embedding,
|
||||
# positions 20-247 use positional_embedding_res.
|
||||
mask1 = torch.zeros(self.CONTEXT_LENGTH, 1)
|
||||
mask1[:20, :] = 1
|
||||
mask2 = torch.zeros(self.CONTEXT_LENGTH, 1)
|
||||
mask2[20:, :] = 1
|
||||
self.register_buffer("mask1", mask1)
|
||||
self.register_buffer("mask2", mask2)
|
||||
|
||||
def _build_attention_mask(self) -> torch.Tensor:
|
||||
mask = torch.empty(self.CONTEXT_LENGTH, self.CONTEXT_LENGTH)
|
||||
mask.fill_(float("-inf"))
|
||||
mask.triu_(1)
|
||||
return mask
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
return self.token_embedding.weight.dtype
|
||||
|
||||
def forward(self, text: torch.Tensor) -> torch.Tensor:
|
||||
"""Encode tokenized text.
|
||||
|
||||
Args:
|
||||
text: Token IDs [B, 248] (int/long).
|
||||
|
||||
Returns:
|
||||
Text embeddings [B, embed_dim].
|
||||
"""
|
||||
x = self.token_embedding(text).type(self.dtype) # [B, 248, D]
|
||||
|
||||
# Dual masked positional embeddings (KPS).
|
||||
x = x + (
|
||||
self.positional_embedding * self.mask1
|
||||
+ self.positional_embedding_res * self.mask2
|
||||
).type(self.dtype)
|
||||
|
||||
x = x.permute(1, 0, 2) # NLD -> LND (sequence-first for nn.MultiheadAttention)
|
||||
x = self.transformer(x)
|
||||
x = x.permute(1, 0, 2) # LND -> NLD
|
||||
x = self.ln_final(x).type(self.dtype)
|
||||
|
||||
# Take features from EOT token (highest token ID in each sequence).
|
||||
x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
|
||||
return x
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Loading and tokenization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _extract_text_state(full_state: dict) -> dict:
|
||||
"""Extract text encoder keys from a full DGTRS-CLIP state dict."""
|
||||
text_keys = {
|
||||
"token_embedding", "positional_embedding", "positional_embedding_res",
|
||||
"transformer", "ln_final", "text_projection",
|
||||
}
|
||||
return {
|
||||
k: v for k, v in full_state.items()
|
||||
if any(k == prefix or k.startswith(prefix + ".") for prefix in text_keys)
|
||||
}
|
||||
|
||||
|
||||
def build_model(state_dict: dict) -> DGTRSTextEncoder:
|
||||
"""Build DGTRSTextEncoder from a DGTRS-CLIP state dict.
|
||||
|
||||
Auto-detects architecture dimensions from the state dict.
|
||||
"""
|
||||
embed_dim = state_dict["text_projection"].shape[1]
|
||||
vocab_size = state_dict["token_embedding.weight"].shape[0]
|
||||
transformer_width = state_dict["ln_final.weight"].shape[0]
|
||||
transformer_heads = transformer_width // 64
|
||||
transformer_layers = len(set(
|
||||
k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")
|
||||
))
|
||||
|
||||
model = DGTRSTextEncoder(
|
||||
vocab_size=vocab_size,
|
||||
transformer_width=transformer_width,
|
||||
transformer_heads=transformer_heads,
|
||||
transformer_layers=transformer_layers,
|
||||
embed_dim=embed_dim,
|
||||
)
|
||||
# mask1/mask2 are buffers created in __init__, not in checkpoint.
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
return model.eval()
|
||||
|
||||
|
||||
def load_dgtrs_text_encoder(
|
||||
checkpoint_path: str | Path,
|
||||
device: str = "cpu",
|
||||
) -> DGTRSTextEncoder:
|
||||
"""Load DGTRS-CLIP text encoder from checkpoint.
|
||||
|
||||
Extracts only text encoder weights from the full CLIP checkpoint.
|
||||
"""
|
||||
path = Path(checkpoint_path)
|
||||
LOGGER.info("📝 Loading DGTRS-CLIP text encoder from %s", path.name)
|
||||
full_state = torch.load(str(path), map_location="cpu", weights_only=False)
|
||||
if "state_dict" in full_state:
|
||||
full_state = full_state["state_dict"]
|
||||
|
||||
text_state = _extract_text_state(full_state)
|
||||
model = build_model(text_state)
|
||||
model = model.to(device)
|
||||
|
||||
n_params = sum(p.numel() for p in model.parameters())
|
||||
LOGGER.info(
|
||||
"📝 DGTRS text encoder loaded: %s params, context=%d tokens",
|
||||
f"{n_params:,}", model.context_length,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def tokenize_dgtrs(
|
||||
texts: str | list[str],
|
||||
context_length: int = 248,
|
||||
truncate: bool = True,
|
||||
) -> torch.Tensor:
|
||||
"""Tokenize text for DGTRS-CLIP (248 token context).
|
||||
|
||||
Args:
|
||||
texts: Input string or list of strings.
|
||||
context_length: Output sequence length (default 248).
|
||||
truncate: Whether to truncate long sequences.
|
||||
|
||||
Returns:
|
||||
Token IDs [B, context_length] (int tensor).
|
||||
"""
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
|
||||
sot_token = _tokenizer.encoder["<|startoftext|>"]
|
||||
eot_token = _tokenizer.encoder["<|endoftext|>"]
|
||||
all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
|
||||
|
||||
result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
|
||||
for i, tokens in enumerate(all_tokens):
|
||||
if len(tokens) > context_length:
|
||||
if truncate:
|
||||
tokens = tokens[:context_length]
|
||||
tokens[-1] = eot_token
|
||||
else:
|
||||
raise RuntimeError(f"Input too long ({len(tokens)} > {context_length})")
|
||||
result[i, :len(tokens)] = torch.tensor(tokens)
|
||||
|
||||
return result
|
||||
125
src/models/dgtrs/simple_tokenizer.py
Normal file
125
src/models/dgtrs/simple_tokenizer.py
Normal file
@@ -0,0 +1,125 @@
|
||||
"""BPE tokenizer for DGTRS-CLIP / LongCLIP.
|
||||
|
||||
Adapted from https://github.com/MitsuiChen14/DGTRS (Apache-2.0).
|
||||
"""
|
||||
|
||||
import gzip
|
||||
import html
|
||||
import os
|
||||
from functools import lru_cache
|
||||
|
||||
import ftfy
|
||||
import regex as re
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def default_bpe():
|
||||
return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def bytes_to_unicode():
|
||||
bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
|
||||
cs = bs[:]
|
||||
n = 0
|
||||
for b in range(2**8):
|
||||
if b not in bs:
|
||||
bs.append(b)
|
||||
cs.append(2**8 + n)
|
||||
n += 1
|
||||
cs = [chr(n) for n in cs]
|
||||
return dict(zip(bs, cs))
|
||||
|
||||
|
||||
def get_pairs(word):
|
||||
pairs = set()
|
||||
prev_char = word[0]
|
||||
for char in word[1:]:
|
||||
pairs.add((prev_char, char))
|
||||
prev_char = char
|
||||
return pairs
|
||||
|
||||
|
||||
def basic_clean(text):
|
||||
text = ftfy.fix_text(text)
|
||||
text = html.unescape(html.unescape(text))
|
||||
return text.strip()
|
||||
|
||||
|
||||
def whitespace_clean(text):
|
||||
text = re.sub(r"\s+", " ", text)
|
||||
text = text.strip()
|
||||
return text
|
||||
|
||||
|
||||
class SimpleTokenizer:
|
||||
def __init__(self, bpe_path: str = default_bpe()):
|
||||
self.byte_encoder = bytes_to_unicode()
|
||||
self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
|
||||
merges = gzip.open(bpe_path).read().decode("utf-8").split("\n")
|
||||
merges = merges[1:49152 - 256 - 2 + 1]
|
||||
merges = [tuple(merge.split()) for merge in merges]
|
||||
vocab = list(bytes_to_unicode().values())
|
||||
vocab = vocab + [v + "</w>" for v in vocab]
|
||||
for merge in merges:
|
||||
vocab.append("".join(merge))
|
||||
vocab.extend(["<|startoftext|>", "<|endoftext|>"])
|
||||
self.encoder = dict(zip(vocab, range(len(vocab))))
|
||||
self.decoder = {v: k for k, v in self.encoder.items()}
|
||||
self.bpe_ranks = dict(zip(merges, range(len(merges))))
|
||||
self.cache = {"<|startoftext|>": "<|startoftext|>", "<|endoftext|>": "<|endoftext|>"}
|
||||
self.pat = re.compile(
|
||||
r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
def bpe(self, token):
|
||||
if token in self.cache:
|
||||
return self.cache[token]
|
||||
word = tuple(token[:-1]) + (token[-1] + "</w>",)
|
||||
pairs = get_pairs(word)
|
||||
if not pairs:
|
||||
return token + "</w>"
|
||||
while True:
|
||||
bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
|
||||
if bigram not in self.bpe_ranks:
|
||||
break
|
||||
first, second = bigram
|
||||
new_word = []
|
||||
i = 0
|
||||
while i < len(word):
|
||||
try:
|
||||
j = word.index(first, i)
|
||||
new_word.extend(word[i:j])
|
||||
i = j
|
||||
except ValueError:
|
||||
new_word.extend(word[i:])
|
||||
break
|
||||
if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
|
||||
new_word.append(first + second)
|
||||
i += 2
|
||||
else:
|
||||
new_word.append(word[i])
|
||||
i += 1
|
||||
new_word = tuple(new_word)
|
||||
word = new_word
|
||||
if len(word) == 1:
|
||||
break
|
||||
else:
|
||||
pairs = get_pairs(word)
|
||||
word = " ".join(word)
|
||||
self.cache[token] = word
|
||||
return word
|
||||
|
||||
def encode(self, text):
|
||||
bpe_tokens = []
|
||||
text = whitespace_clean(basic_clean(text)).lower()
|
||||
for token in re.findall(self.pat, text):
|
||||
token = "".join(self.byte_encoder[b] for b in token.encode("utf-8"))
|
||||
bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(" "))
|
||||
return bpe_tokens
|
||||
|
||||
def decode(self, tokens):
|
||||
text = "".join([self.decoder[token] for token in tokens])
|
||||
text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors="replace").replace("</w>", " ")
|
||||
return text
|
||||
Reference in New Issue
Block a user