Suya OCR API — vLLM-backed, OpenAI-compatible OCR service
FastAPI service wrapping the Surya-OCR-2 model (datalab-to) served through vLLM: legacy /v1/api/ai/* endpoints, an OpenAI-compatible /v1/chat/completions endpoint, a coalescing request batcher, a local OCR CLI, Docker packaging, multilingual example outputs, and quantization/concurrency benchmarks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Fu Dai
@@ -0,0 +1,52 @@
|
||||
import math
|
||||
from typing import List, Optional
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from surya.common.predictor import BasePredictor
|
||||
from surya.ocr_error.loader import OCRErrorModelLoader
|
||||
from surya.ocr_error.model.config import ID2LABEL
|
||||
from surya.ocr_error.schema import OCRErrorDetectionResult
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
class OCRErrorPredictor(BasePredictor):
|
||||
model_loader_cls = OCRErrorModelLoader
|
||||
batch_size = settings.OCR_ERROR_BATCH_SIZE
|
||||
default_batch_sizes = {"cpu": 8, "mps": 8, "cuda": 64}
|
||||
|
||||
def __call__(self, texts: List[str], batch_size: Optional[int] = None):
|
||||
return self.batch_ocr_error_detection(texts, batch_size)
|
||||
|
||||
def batch_ocr_error_detection(
|
||||
self, texts: List[str], batch_size: Optional[int] = None
|
||||
):
|
||||
if batch_size is None:
|
||||
batch_size = self.get_batch_size()
|
||||
|
||||
num_batches = math.ceil(len(texts) / batch_size)
|
||||
texts_processed = self.processor(
|
||||
texts, padding="longest", truncation=True, return_tensors="pt"
|
||||
)
|
||||
predictions = []
|
||||
for batch_idx in tqdm(
|
||||
range(num_batches),
|
||||
desc="Running OCR Error Detection",
|
||||
disable=self.disable_tqdm,
|
||||
):
|
||||
start_idx, end_idx = batch_idx * batch_size, (batch_idx + 1) * batch_size
|
||||
batch_input_ids = texts_processed.input_ids[start_idx:end_idx].to(
|
||||
self.model.device
|
||||
)
|
||||
batch_attention_mask = texts_processed.attention_mask[start_idx:end_idx].to(
|
||||
self.model.device
|
||||
)
|
||||
|
||||
with settings.INFERENCE_MODE():
|
||||
pred = self.model(batch_input_ids, attention_mask=batch_attention_mask)
|
||||
logits = pred.logits.argmax(dim=1).cpu().tolist()
|
||||
predictions.extend(logits)
|
||||
|
||||
return OCRErrorDetectionResult(
|
||||
texts=texts, labels=[ID2LABEL[p] for p in predictions]
|
||||
)
|
||||
@@ -0,0 +1,48 @@
|
||||
from typing import Optional
|
||||
|
||||
|
||||
from surya.common.load import ModelLoader
|
||||
from surya.logging import get_logger
|
||||
from surya.ocr_error.model.config import DistilBertConfig
|
||||
from surya.ocr_error.model.encoder import DistilBertForSequenceClassification
|
||||
from surya.ocr_error.tokenizer import DistilBertTokenizer
|
||||
from surya.settings import settings
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class OCRErrorModelLoader(ModelLoader):
|
||||
def __init__(self, checkpoint: Optional[str] = None):
|
||||
super().__init__(checkpoint)
|
||||
|
||||
if self.checkpoint is None:
|
||||
self.checkpoint = settings.OCR_ERROR_MODEL_CHECKPOINT
|
||||
|
||||
def model(
|
||||
self,
|
||||
device=settings.TORCH_DEVICE_MODEL,
|
||||
dtype=settings.MODEL_DTYPE,
|
||||
attention_implementation: Optional[str] = None,
|
||||
) -> DistilBertForSequenceClassification:
|
||||
if device is None:
|
||||
device = settings.TORCH_DEVICE_MODEL
|
||||
if dtype is None:
|
||||
dtype = settings.MODEL_DTYPE
|
||||
|
||||
config = DistilBertConfig.from_pretrained(self.checkpoint)
|
||||
model = (
|
||||
DistilBertForSequenceClassification.from_pretrained(
|
||||
self.checkpoint,
|
||||
dtype=dtype,
|
||||
config=config,
|
||||
)
|
||||
.to(device)
|
||||
.eval()
|
||||
)
|
||||
|
||||
return model
|
||||
|
||||
def processor(
|
||||
self, device=settings.TORCH_DEVICE_MODEL, dtype=settings.MODEL_DTYPE
|
||||
) -> DistilBertTokenizer:
|
||||
return DistilBertTokenizer.from_pretrained(self.checkpoint)
|
||||
@@ -0,0 +1,68 @@
|
||||
from collections import OrderedDict
|
||||
from typing import Mapping
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.onnx import OnnxConfig
|
||||
|
||||
from surya.common.s3 import S3DownloaderMixin
|
||||
|
||||
ID2LABEL = {
|
||||
0: 'good',
|
||||
1: 'bad'
|
||||
}
|
||||
|
||||
class DistilBertConfig(S3DownloaderMixin, PretrainedConfig):
|
||||
model_type = "distilbert"
|
||||
attribute_map = {
|
||||
"hidden_size": "dim",
|
||||
"num_attention_heads": "n_heads",
|
||||
"num_hidden_layers": "n_layers",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=30522,
|
||||
max_position_embeddings=512,
|
||||
sinusoidal_pos_embds=False,
|
||||
n_layers=6,
|
||||
n_heads=12,
|
||||
dim=768,
|
||||
hidden_dim=4 * 768,
|
||||
dropout=0.1,
|
||||
attention_dropout=0.1,
|
||||
activation="gelu",
|
||||
initializer_range=0.02,
|
||||
qa_dropout=0.1,
|
||||
seq_classif_dropout=0.2,
|
||||
pad_token_id=0,
|
||||
**kwargs,
|
||||
):
|
||||
self.vocab_size = vocab_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.sinusoidal_pos_embds = sinusoidal_pos_embds
|
||||
self.n_layers = n_layers
|
||||
self.n_heads = n_heads
|
||||
self.dim = dim
|
||||
self.hidden_dim = hidden_dim
|
||||
self.dropout = dropout
|
||||
self.attention_dropout = attention_dropout
|
||||
self.activation = activation
|
||||
self.initializer_range = initializer_range
|
||||
self.qa_dropout = qa_dropout
|
||||
self.seq_classif_dropout = seq_classif_dropout
|
||||
super().__init__(**kwargs, pad_token_id=pad_token_id)
|
||||
|
||||
|
||||
class DistilBertOnnxConfig(OnnxConfig):
|
||||
@property
|
||||
def inputs(self) -> Mapping[str, Mapping[int, str]]:
|
||||
if self.task == "multiple-choice":
|
||||
dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
|
||||
else:
|
||||
dynamic_axis = {0: "batch", 1: "sequence"}
|
||||
return OrderedDict(
|
||||
[
|
||||
("input_ids", dynamic_axis),
|
||||
("attention_mask", dynamic_axis),
|
||||
]
|
||||
)
|
||||
@@ -0,0 +1,910 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Optional, Set, List, Tuple, Union, Dict
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F, MSELoss, CrossEntropyLoss, BCEWithLogitsLoss
|
||||
from transformers import apply_chunking_to_forward
|
||||
from transformers.activations import get_activation
|
||||
from transformers.modeling_outputs import BaseModelOutput, SequenceClassifierOutput
|
||||
from transformers.pytorch_utils import (
|
||||
find_pruneable_heads_and_indices,
|
||||
prune_linear_layer,
|
||||
)
|
||||
|
||||
from transformers.utils import (
|
||||
is_flash_attn_greater_or_equal_2_10,
|
||||
)
|
||||
|
||||
from surya.common.pretrained import SuryaPreTrainedModel
|
||||
|
||||
from surya.common.s3 import S3DownloaderMixin
|
||||
from surya.ocr_error.model.config import DistilBertConfig
|
||||
|
||||
|
||||
def _get_unpad_data(attention_mask):
|
||||
seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
|
||||
indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
|
||||
max_seqlen_in_batch = seqlens_in_batch.max().item()
|
||||
cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
|
||||
return (
|
||||
indices,
|
||||
cu_seqlens,
|
||||
max_seqlen_in_batch,
|
||||
)
|
||||
|
||||
|
||||
def create_sinusoidal_embeddings(n_pos: int, dim: int, out: torch.Tensor):
|
||||
position_enc = np.array(
|
||||
[
|
||||
[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)]
|
||||
for pos in range(n_pos)
|
||||
]
|
||||
)
|
||||
out.requires_grad = False
|
||||
out[:, 0::2] = torch.FloatTensor(np.sin(position_enc[:, 0::2]))
|
||||
out[:, 1::2] = torch.FloatTensor(np.cos(position_enc[:, 1::2]))
|
||||
out.detach_()
|
||||
|
||||
|
||||
class Embeddings(nn.Module):
|
||||
def __init__(self, config: DistilBertConfig):
|
||||
super().__init__()
|
||||
self.word_embeddings = nn.Embedding(
|
||||
config.vocab_size, config.dim, padding_idx=config.pad_token_id
|
||||
)
|
||||
self.position_embeddings = nn.Embedding(
|
||||
config.max_position_embeddings, config.dim
|
||||
)
|
||||
|
||||
self.LayerNorm = nn.LayerNorm(config.dim, eps=1e-12)
|
||||
self.dropout = nn.Dropout(config.dropout)
|
||||
self.register_buffer(
|
||||
"position_ids",
|
||||
torch.arange(config.max_position_embeddings).expand((1, -1)),
|
||||
persistent=False,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, input_ids: torch.Tensor, input_embeds: Optional[torch.Tensor] = None
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Parameters:
|
||||
input_ids (torch.Tensor):
|
||||
torch.tensor(bs, max_seq_length) The token ids to embed.
|
||||
input_embeds (*optional*, torch.Tensor):
|
||||
The pre-computed word embeddings. Can only be passed if the input ids are `None`.
|
||||
|
||||
|
||||
Returns: torch.tensor(bs, max_seq_length, dim) The embedded tokens (plus position embeddings, no token_type
|
||||
embeddings)
|
||||
"""
|
||||
if input_ids is not None:
|
||||
input_embeds = self.word_embeddings(input_ids) # (bs, max_seq_length, dim)
|
||||
|
||||
seq_length = input_embeds.size(1)
|
||||
|
||||
# Setting the position-ids to the registered buffer in constructor, it helps
|
||||
# when tracing the model without passing position-ids, solves
|
||||
# isues similar to issue #5664
|
||||
if hasattr(self, "position_ids"):
|
||||
position_ids = self.position_ids[:, :seq_length]
|
||||
else:
|
||||
position_ids = torch.arange(
|
||||
seq_length, dtype=torch.long, device=input_ids.device
|
||||
) # (max_seq_length)
|
||||
position_ids = position_ids.unsqueeze(0).expand_as(
|
||||
input_ids
|
||||
) # (bs, max_seq_length)
|
||||
|
||||
position_embeddings = self.position_embeddings(
|
||||
position_ids
|
||||
) # (bs, max_seq_length, dim)
|
||||
|
||||
embeddings = input_embeds + position_embeddings # (bs, max_seq_length, dim)
|
||||
embeddings = self.LayerNorm(embeddings) # (bs, max_seq_length, dim)
|
||||
embeddings = self.dropout(embeddings) # (bs, max_seq_length, dim)
|
||||
return embeddings
|
||||
|
||||
|
||||
class MultiHeadSelfAttention(nn.Module):
|
||||
def __init__(self, config: DistilBertConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
self.n_heads = config.n_heads
|
||||
self.dim = config.dim
|
||||
self.dropout = nn.Dropout(p=config.attention_dropout)
|
||||
self.is_causal = False
|
||||
|
||||
# Have an even number of multi heads that divide the dimensions
|
||||
if self.dim % self.n_heads != 0:
|
||||
# Raise value errors for even multi-head attention nodes
|
||||
raise ValueError(
|
||||
f"self.n_heads: {self.n_heads} must divide self.dim: {self.dim} evenly"
|
||||
)
|
||||
|
||||
self.q_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
|
||||
self.k_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
|
||||
self.v_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
|
||||
self.out_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
|
||||
|
||||
self.pruned_heads: Set[int] = set()
|
||||
self.attention_head_size = self.dim // self.n_heads
|
||||
|
||||
def prune_heads(self, heads: List[int]):
|
||||
if len(heads) == 0:
|
||||
return
|
||||
heads, index = find_pruneable_heads_and_indices(
|
||||
heads, self.n_heads, self.attention_head_size, self.pruned_heads
|
||||
)
|
||||
# Prune linear layers
|
||||
self.q_lin = prune_linear_layer(self.q_lin, index)
|
||||
self.k_lin = prune_linear_layer(self.k_lin, index)
|
||||
self.v_lin = prune_linear_layer(self.v_lin, index)
|
||||
self.out_lin = prune_linear_layer(self.out_lin, index, dim=1)
|
||||
# Update hyper params
|
||||
self.n_heads = self.n_heads - len(heads)
|
||||
self.dim = self.attention_head_size * self.n_heads
|
||||
self.pruned_heads = self.pruned_heads.union(heads)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
head_mask: Optional[torch.Tensor] = None,
|
||||
output_attentions: bool = False,
|
||||
) -> Tuple[torch.Tensor, ...]:
|
||||
"""
|
||||
Parameters:
|
||||
query: torch.tensor(bs, seq_length, dim)
|
||||
key: torch.tensor(bs, seq_length, dim)
|
||||
value: torch.tensor(bs, seq_length, dim)
|
||||
mask: torch.tensor(bs, seq_length)
|
||||
|
||||
Returns:
|
||||
weights: torch.tensor(bs, n_heads, seq_length, seq_length) Attention weights context: torch.tensor(bs,
|
||||
seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True`
|
||||
"""
|
||||
bs, q_length, dim = query.size()
|
||||
k_length = key.size(1)
|
||||
# assert dim == self.dim, f'Dimensions do not match: {dim} input vs {self.dim} configured'
|
||||
# assert key.size() == value.size()
|
||||
|
||||
dim_per_head = self.dim // self.n_heads
|
||||
|
||||
mask_reshp = (bs, 1, 1, k_length)
|
||||
|
||||
def shape(x: torch.Tensor) -> torch.Tensor:
|
||||
"""separate heads"""
|
||||
return x.view(bs, -1, self.n_heads, dim_per_head).transpose(1, 2)
|
||||
|
||||
def unshape(x: torch.Tensor) -> torch.Tensor:
|
||||
"""group heads"""
|
||||
return (
|
||||
x.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * dim_per_head)
|
||||
)
|
||||
|
||||
q = shape(self.q_lin(query)) # (bs, n_heads, q_length, dim_per_head)
|
||||
k = shape(self.k_lin(key)) # (bs, n_heads, k_length, dim_per_head)
|
||||
v = shape(self.v_lin(value)) # (bs, n_heads, k_length, dim_per_head)
|
||||
|
||||
q = q / math.sqrt(dim_per_head) # (bs, n_heads, q_length, dim_per_head)
|
||||
scores = torch.matmul(q, k.transpose(2, 3)) # (bs, n_heads, q_length, k_length)
|
||||
mask = (
|
||||
(mask == 0).view(mask_reshp).expand_as(scores)
|
||||
) # (bs, n_heads, q_length, k_length)
|
||||
scores = scores.masked_fill(
|
||||
mask, torch.tensor(torch.finfo(scores.dtype).min)
|
||||
) # (bs, n_heads, q_length, k_length)
|
||||
|
||||
weights = nn.functional.softmax(
|
||||
scores, dim=-1
|
||||
) # (bs, n_heads, q_length, k_length)
|
||||
weights = self.dropout(weights) # (bs, n_heads, q_length, k_length)
|
||||
|
||||
# Mask heads if we want to
|
||||
if head_mask is not None:
|
||||
weights = weights * head_mask
|
||||
|
||||
context = torch.matmul(weights, v) # (bs, n_heads, q_length, dim_per_head)
|
||||
context = unshape(context) # (bs, q_length, dim)
|
||||
context = self.out_lin(context) # (bs, q_length, dim)
|
||||
|
||||
if output_attentions:
|
||||
return (context, weights)
|
||||
else:
|
||||
return (context,)
|
||||
|
||||
|
||||
class DistilBertFlashAttention2(MultiHeadSelfAttention):
|
||||
"""
|
||||
DistilBert flash attention module. This module inherits from `MultiHeadSelfAttention` as the weights of the module
|
||||
stays untouched. The only required change would be on the forward pass where it needs to correctly call the public
|
||||
API of flash attention and deal with padding tokens in case the input contains any of them.
|
||||
"""
|
||||
|
||||
# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
# TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
|
||||
# flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
|
||||
# Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
|
||||
self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
head_mask: Optional[torch.Tensor] = None,
|
||||
output_attentions: bool = False,
|
||||
) -> Tuple[torch.Tensor, ...]:
|
||||
"""
|
||||
Parameters:
|
||||
query: torch.tensor(bs, seq_length, dim)
|
||||
key: torch.tensor(bs, seq_length, dim)
|
||||
value: torch.tensor(bs, seq_length, dim)
|
||||
mask: torch.tensor(bs, seq_length)
|
||||
|
||||
Returns:
|
||||
weights: torch.tensor(bs, n_heads, seq_length, seq_length) Attention weights context: torch.tensor(bs,
|
||||
seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True`
|
||||
"""
|
||||
batch_size, q_length, dim = query.size()
|
||||
|
||||
dim_per_head = self.dim // self.n_heads
|
||||
|
||||
def reshape(x: torch.Tensor) -> torch.Tensor:
|
||||
"""separate heads"""
|
||||
return x.view(batch_size, -1, self.n_heads, dim_per_head)
|
||||
|
||||
# Flash attention requires the input to have the shape
|
||||
# batch_size x seq_length x head_dim x hidden_dim
|
||||
query_states = reshape(self.q_lin(query))
|
||||
key_states = reshape(self.k_lin(key))
|
||||
value_states = reshape(self.v_lin(value))
|
||||
|
||||
attn_dropout = self.config.attention_dropout if self.training else 0.0
|
||||
|
||||
# In PEFT, usually we cast the layer norms in float32 for training stability reasons
|
||||
# therefore the input hidden states gets silently casted in float32. Hence, we need
|
||||
# cast them back in the correct dtype just to be sure everything works as expected.
|
||||
# This might slowdown training & inference so it is recommended to not cast the LayerNorms
|
||||
# in fp32. (LlamaRMSNorm handles it correctly)
|
||||
|
||||
if query_states.dtype == torch.float32:
|
||||
if torch.is_autocast_enabled():
|
||||
target_dtype = torch.get_autocast_gpu_dtype()
|
||||
# Handle the case where the model is quantized
|
||||
elif hasattr(self.config, "_pre_quantization_dtype"):
|
||||
target_dtype = self.config._pre_quantization_dtype
|
||||
else:
|
||||
target_dtype = self.q_lin.weight.dtype
|
||||
|
||||
query_states = query_states.to(target_dtype)
|
||||
key_states = key_states.to(target_dtype)
|
||||
value_states = value_states.to(target_dtype)
|
||||
|
||||
attn_weights = self._flash_attention_forward(
|
||||
query_states, key_states, value_states, mask, q_length, dropout=attn_dropout
|
||||
)
|
||||
|
||||
attn_weights_reshaped = attn_weights.reshape(
|
||||
batch_size, q_length, self.n_heads * dim_per_head
|
||||
)
|
||||
attn_output = self.out_lin(attn_weights_reshaped)
|
||||
|
||||
if output_attentions:
|
||||
return (attn_output, attn_weights)
|
||||
else:
|
||||
return (attn_output,)
|
||||
|
||||
# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward with causal=True->causal=False
|
||||
def _flash_attention_forward(
|
||||
self,
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
attention_mask,
|
||||
query_length,
|
||||
dropout=0.0,
|
||||
softmax_scale=None,
|
||||
):
|
||||
"""
|
||||
Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
|
||||
first unpad the input, then computes the attention scores and pad the final attention scores.
|
||||
|
||||
Args:
|
||||
query_states (`torch.Tensor`):
|
||||
Input query states to be passed to Flash Attention API
|
||||
key_states (`torch.Tensor`):
|
||||
Input key states to be passed to Flash Attention API
|
||||
value_states (`torch.Tensor`):
|
||||
Input value states to be passed to Flash Attention API
|
||||
attention_mask (`torch.Tensor`):
|
||||
The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
|
||||
position of padding tokens and 1 for the position of non-padding tokens.
|
||||
dropout (`float`):
|
||||
Attention dropout
|
||||
softmax_scale (`float`, *optional*):
|
||||
The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
|
||||
"""
|
||||
from flash_attn import flash_attn_func, flash_attn_varlen_func
|
||||
from flash_attn.bert_padding import pad_input
|
||||
|
||||
if not self._flash_attn_uses_top_left_mask:
|
||||
causal = self.is_causal
|
||||
else:
|
||||
# TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
|
||||
causal = self.is_causal and query_length != 1
|
||||
|
||||
# Contains at least one padding token in the sequence
|
||||
if attention_mask is not None:
|
||||
batch_size = query_states.shape[0]
|
||||
(
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
indices_q,
|
||||
cu_seq_lens,
|
||||
max_seq_lens,
|
||||
) = self._upad_input(
|
||||
query_states, key_states, value_states, attention_mask, query_length
|
||||
)
|
||||
|
||||
cu_seqlens_q, cu_seqlens_k = cu_seq_lens
|
||||
max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
|
||||
|
||||
attn_output_unpad = flash_attn_varlen_func(
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
max_seqlen_q=max_seqlen_in_batch_q,
|
||||
max_seqlen_k=max_seqlen_in_batch_k,
|
||||
dropout_p=dropout,
|
||||
softmax_scale=softmax_scale,
|
||||
causal=causal,
|
||||
)
|
||||
|
||||
attn_output = pad_input(
|
||||
attn_output_unpad, indices_q, batch_size, query_length
|
||||
)
|
||||
else:
|
||||
attn_output = flash_attn_func(
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
dropout,
|
||||
softmax_scale=softmax_scale,
|
||||
causal=causal,
|
||||
)
|
||||
|
||||
return attn_output
|
||||
|
||||
# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input with num_heads->n_heads
|
||||
def _upad_input(
|
||||
self, query_layer, key_layer, value_layer, attention_mask, query_length
|
||||
):
|
||||
from flash_attn.bert_padding import index_first_axis, unpad_input
|
||||
|
||||
indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
|
||||
batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
|
||||
|
||||
key_layer = index_first_axis(
|
||||
key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
|
||||
indices_k,
|
||||
)
|
||||
value_layer = index_first_axis(
|
||||
value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
|
||||
indices_k,
|
||||
)
|
||||
if query_length == kv_seq_len:
|
||||
query_layer = index_first_axis(
|
||||
query_layer.reshape(batch_size * kv_seq_len, self.n_heads, head_dim),
|
||||
indices_k,
|
||||
)
|
||||
cu_seqlens_q = cu_seqlens_k
|
||||
max_seqlen_in_batch_q = max_seqlen_in_batch_k
|
||||
indices_q = indices_k
|
||||
elif query_length == 1:
|
||||
max_seqlen_in_batch_q = 1
|
||||
cu_seqlens_q = torch.arange(
|
||||
batch_size + 1, dtype=torch.int32, device=query_layer.device
|
||||
) # There is a memcpy here, that is very bad.
|
||||
indices_q = cu_seqlens_q[:-1]
|
||||
query_layer = query_layer.squeeze(1)
|
||||
else:
|
||||
# The -q_len: slice assumes left padding.
|
||||
attention_mask = attention_mask[:, -query_length:]
|
||||
query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
|
||||
query_layer, attention_mask
|
||||
)
|
||||
|
||||
return (
|
||||
query_layer,
|
||||
key_layer,
|
||||
value_layer,
|
||||
indices_q,
|
||||
(cu_seqlens_q, cu_seqlens_k),
|
||||
(max_seqlen_in_batch_q, max_seqlen_in_batch_k),
|
||||
)
|
||||
|
||||
|
||||
class FFN(nn.Module):
|
||||
def __init__(self, config: DistilBertConfig):
|
||||
super().__init__()
|
||||
self.dropout = nn.Dropout(p=config.dropout)
|
||||
self.chunk_size_feed_forward = config.chunk_size_feed_forward
|
||||
self.seq_len_dim = 1
|
||||
self.lin1 = nn.Linear(in_features=config.dim, out_features=config.hidden_dim)
|
||||
self.lin2 = nn.Linear(in_features=config.hidden_dim, out_features=config.dim)
|
||||
self.activation = get_activation(config.activation)
|
||||
|
||||
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
||||
return apply_chunking_to_forward(
|
||||
self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, input
|
||||
)
|
||||
|
||||
def ff_chunk(self, input: torch.Tensor) -> torch.Tensor:
|
||||
x = self.lin1(input)
|
||||
x = self.activation(x)
|
||||
x = self.lin2(x)
|
||||
x = self.dropout(x)
|
||||
return x
|
||||
|
||||
|
||||
DISTILBERT_ATTENTION_CLASSES = {
|
||||
"eager": MultiHeadSelfAttention,
|
||||
"flash_attention_2": DistilBertFlashAttention2,
|
||||
}
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, config: DistilBertConfig):
|
||||
super().__init__()
|
||||
|
||||
# Have an even number of Configure multi-heads
|
||||
if config.dim % config.n_heads != 0:
|
||||
raise ValueError(
|
||||
f"config.n_heads {config.n_heads} must divide config.dim {config.dim} evenly"
|
||||
)
|
||||
|
||||
self.attention = DISTILBERT_ATTENTION_CLASSES[config._attn_implementation](
|
||||
config
|
||||
)
|
||||
self.sa_layer_norm = nn.LayerNorm(normalized_shape=config.dim, eps=1e-12)
|
||||
|
||||
self.ffn = FFN(config)
|
||||
self.output_layer_norm = nn.LayerNorm(normalized_shape=config.dim, eps=1e-12)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
attn_mask: Optional[torch.Tensor] = None,
|
||||
head_mask: Optional[torch.Tensor] = None,
|
||||
output_attentions: bool = False,
|
||||
) -> Tuple[torch.Tensor, ...]:
|
||||
"""
|
||||
Parameters:
|
||||
x: torch.tensor(bs, seq_length, dim)
|
||||
attn_mask: torch.tensor(bs, seq_length)
|
||||
|
||||
Returns:
|
||||
sa_weights: torch.tensor(bs, n_heads, seq_length, seq_length) The attention weights ffn_output:
|
||||
torch.tensor(bs, seq_length, dim) The output of the transformer block contextualization.
|
||||
"""
|
||||
# Self-Attention
|
||||
sa_output = self.attention(
|
||||
query=x,
|
||||
key=x,
|
||||
value=x,
|
||||
mask=attn_mask,
|
||||
head_mask=head_mask,
|
||||
output_attentions=output_attentions,
|
||||
)
|
||||
if output_attentions:
|
||||
sa_output, sa_weights = (
|
||||
sa_output # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length)
|
||||
)
|
||||
else: # To handle these `output_attentions` or `output_hidden_states` cases returning tuples
|
||||
sa_output = sa_output[0]
|
||||
|
||||
sa_output = self.sa_layer_norm(sa_output + x) # (bs, seq_length, dim)
|
||||
|
||||
# Feed Forward Network
|
||||
ffn_output = self.ffn(sa_output) # (bs, seq_length, dim)
|
||||
ffn_output: torch.Tensor = self.output_layer_norm(
|
||||
ffn_output + sa_output
|
||||
) # (bs, seq_length, dim)
|
||||
|
||||
output = (ffn_output,)
|
||||
if output_attentions:
|
||||
output = (sa_weights,) + output
|
||||
return output
|
||||
|
||||
|
||||
class Transformer(nn.Module):
|
||||
def __init__(self, config: DistilBertConfig):
|
||||
super().__init__()
|
||||
self.n_layers = config.n_layers
|
||||
self.layer = nn.ModuleList(
|
||||
[TransformerBlock(config) for _ in range(config.n_layers)]
|
||||
)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
attn_mask: Optional[torch.Tensor] = None,
|
||||
head_mask: Optional[torch.Tensor] = None,
|
||||
output_attentions: bool = False,
|
||||
output_hidden_states: bool = False,
|
||||
return_dict: Optional[bool] = None,
|
||||
) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]: # docstyle-ignore
|
||||
"""
|
||||
Parameters:
|
||||
x: torch.tensor(bs, seq_length, dim) Input sequence embedded.
|
||||
attn_mask: torch.tensor(bs, seq_length) Attention mask on the sequence.
|
||||
|
||||
Returns:
|
||||
hidden_state: torch.tensor(bs, seq_length, dim) Sequence of hidden states in the last (top)
|
||||
layer all_hidden_states: Tuple[torch.tensor(bs, seq_length, dim)]
|
||||
Tuple of length n_layers with the hidden states from each layer.
|
||||
Optional: only if output_hidden_states=True
|
||||
all_attentions: Tuple[torch.tensor(bs, n_heads, seq_length, seq_length)]
|
||||
Tuple of length n_layers with the attention weights from each layer
|
||||
Optional: only if output_attentions=True
|
||||
"""
|
||||
all_hidden_states = () if output_hidden_states else None
|
||||
all_attentions = () if output_attentions else None
|
||||
|
||||
hidden_state = x
|
||||
for i, layer_module in enumerate(self.layer):
|
||||
if output_hidden_states:
|
||||
all_hidden_states = all_hidden_states + (hidden_state,)
|
||||
|
||||
if self.gradient_checkpointing and self.training:
|
||||
layer_outputs = self._gradient_checkpointing_func(
|
||||
layer_module.__call__,
|
||||
hidden_state,
|
||||
attn_mask,
|
||||
head_mask[i],
|
||||
output_attentions,
|
||||
)
|
||||
else:
|
||||
layer_outputs = layer_module(
|
||||
hidden_state,
|
||||
attn_mask,
|
||||
head_mask[i],
|
||||
output_attentions,
|
||||
)
|
||||
|
||||
hidden_state = layer_outputs[-1]
|
||||
|
||||
if output_attentions:
|
||||
if len(layer_outputs) != 2:
|
||||
raise ValueError(
|
||||
f"The length of the layer_outputs should be 2, but it is {len(layer_outputs)}"
|
||||
)
|
||||
|
||||
attentions = layer_outputs[0]
|
||||
all_attentions = all_attentions + (attentions,)
|
||||
else:
|
||||
if len(layer_outputs) != 1:
|
||||
raise ValueError(
|
||||
f"The length of the layer_outputs should be 1, but it is {len(layer_outputs)}"
|
||||
)
|
||||
|
||||
# Add last layer
|
||||
if output_hidden_states:
|
||||
all_hidden_states = all_hidden_states + (hidden_state,)
|
||||
|
||||
if not return_dict:
|
||||
return tuple(
|
||||
v
|
||||
for v in [hidden_state, all_hidden_states, all_attentions]
|
||||
if v is not None
|
||||
)
|
||||
return BaseModelOutput(
|
||||
last_hidden_state=hidden_state,
|
||||
hidden_states=all_hidden_states,
|
||||
attentions=all_attentions,
|
||||
)
|
||||
|
||||
|
||||
class DistilBertPreTrainedModel(SuryaPreTrainedModel):
|
||||
"""
|
||||
An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
|
||||
models.
|
||||
"""
|
||||
|
||||
config_class = DistilBertConfig
|
||||
load_tf_weights = None
|
||||
base_model_prefix = "distilbert"
|
||||
supports_gradient_checkpointing = True
|
||||
_supports_flash_attn_2 = True
|
||||
|
||||
def _init_weights(self, module: nn.Module):
|
||||
"""Initialize the weights."""
|
||||
if isinstance(module, nn.Linear):
|
||||
# Slightly different from the TF version which uses truncated_normal for initialization
|
||||
# cf https://github.com/pytorch/pytorch/pull/5617
|
||||
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
|
||||
if module.bias is not None:
|
||||
module.bias.data.zero_()
|
||||
elif isinstance(module, nn.Embedding):
|
||||
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
|
||||
if module.padding_idx is not None:
|
||||
module.weight.data[module.padding_idx].zero_()
|
||||
elif isinstance(module, nn.LayerNorm):
|
||||
module.bias.data.zero_()
|
||||
module.weight.data.fill_(1.0)
|
||||
elif isinstance(module, Embeddings) and self.config.sinusoidal_pos_embds:
|
||||
create_sinusoidal_embeddings(
|
||||
self.config.max_position_embeddings,
|
||||
self.config.dim,
|
||||
module.position_embeddings.weight,
|
||||
)
|
||||
|
||||
|
||||
class DistilBertModel(DistilBertPreTrainedModel):
|
||||
def __init__(self, config: DistilBertConfig):
|
||||
super().__init__(config)
|
||||
|
||||
self.embeddings = Embeddings(config) # Embeddings
|
||||
self.transformer = Transformer(config) # Encoder
|
||||
self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def get_position_embeddings(self) -> nn.Embedding:
|
||||
"""
|
||||
Returns the position embeddings
|
||||
"""
|
||||
return self.embeddings.position_embeddings
|
||||
|
||||
def resize_position_embeddings(self, new_num_position_embeddings: int):
|
||||
"""
|
||||
Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.
|
||||
|
||||
Arguments:
|
||||
new_num_position_embeddings (`int`):
|
||||
The number of new position embedding matrix. If position embeddings are learned, increasing the size
|
||||
will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the
|
||||
end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the
|
||||
size will add correct vectors at the end following the position encoding algorithm, whereas reducing
|
||||
the size will remove vectors from the end.
|
||||
"""
|
||||
num_position_embeds_diff = (
|
||||
new_num_position_embeddings - self.config.max_position_embeddings
|
||||
)
|
||||
|
||||
# no resizing needs to be done if the length stays the same
|
||||
if num_position_embeds_diff == 0:
|
||||
return
|
||||
|
||||
self.config.max_position_embeddings = new_num_position_embeddings
|
||||
|
||||
old_position_embeddings_weight = (
|
||||
self.embeddings.position_embeddings.weight.clone()
|
||||
)
|
||||
|
||||
self.embeddings.position_embeddings = nn.Embedding(
|
||||
self.config.max_position_embeddings, self.config.dim
|
||||
)
|
||||
|
||||
if self.config.sinusoidal_pos_embds:
|
||||
create_sinusoidal_embeddings(
|
||||
n_pos=self.config.max_position_embeddings,
|
||||
dim=self.config.dim,
|
||||
out=self.position_embeddings.weight,
|
||||
)
|
||||
else:
|
||||
with torch.no_grad():
|
||||
if num_position_embeds_diff > 0:
|
||||
self.embeddings.position_embeddings.weight[
|
||||
:-num_position_embeds_diff
|
||||
] = nn.Parameter(old_position_embeddings_weight)
|
||||
else:
|
||||
self.embeddings.position_embeddings.weight = nn.Parameter(
|
||||
old_position_embeddings_weight[:num_position_embeds_diff]
|
||||
)
|
||||
# move position_embeddings to correct device
|
||||
self.embeddings.position_embeddings.to(self.device)
|
||||
|
||||
def get_input_embeddings(self) -> nn.Embedding:
|
||||
return self.embeddings.word_embeddings
|
||||
|
||||
def set_input_embeddings(self, new_embeddings: nn.Embedding):
|
||||
self.embeddings.word_embeddings = new_embeddings
|
||||
|
||||
def _prune_heads(self, heads_to_prune: Dict[int, List[List[int]]]):
|
||||
"""
|
||||
Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
|
||||
class PreTrainedModel
|
||||
"""
|
||||
for layer, heads in heads_to_prune.items():
|
||||
self.transformer.layer[layer].attention.prune_heads(heads)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
head_mask: Optional[torch.Tensor] = None,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]:
|
||||
output_attentions = (
|
||||
output_attentions
|
||||
if output_attentions is not None
|
||||
else self.config.output_attentions
|
||||
)
|
||||
output_hidden_states = (
|
||||
output_hidden_states
|
||||
if output_hidden_states is not None
|
||||
else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = (
|
||||
return_dict if return_dict is not None else self.config.use_return_dict
|
||||
)
|
||||
|
||||
if input_ids is not None and inputs_embeds is not None:
|
||||
raise ValueError(
|
||||
"You cannot specify both input_ids and inputs_embeds at the same time"
|
||||
)
|
||||
elif input_ids is not None:
|
||||
self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
|
||||
input_shape = input_ids.size()
|
||||
elif inputs_embeds is not None:
|
||||
input_shape = inputs_embeds.size()[:-1]
|
||||
else:
|
||||
raise ValueError("You have to specify either input_ids or inputs_embeds")
|
||||
|
||||
device = input_ids.device if input_ids is not None else inputs_embeds.device
|
||||
|
||||
# Prepare head mask if needed
|
||||
head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
|
||||
|
||||
embeddings = self.embeddings(input_ids, inputs_embeds) # (bs, seq_length, dim)
|
||||
|
||||
if self._use_flash_attention_2:
|
||||
attention_mask = (
|
||||
attention_mask
|
||||
if (attention_mask is not None and 0 in attention_mask)
|
||||
else None
|
||||
)
|
||||
else:
|
||||
if attention_mask is None:
|
||||
attention_mask = torch.ones(
|
||||
input_shape, device=device
|
||||
) # (bs, seq_length)
|
||||
|
||||
return self.transformer(
|
||||
x=embeddings,
|
||||
attn_mask=attention_mask,
|
||||
head_mask=head_mask,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
|
||||
|
||||
class DistilBertForSequenceClassification(S3DownloaderMixin, DistilBertPreTrainedModel):
|
||||
def __init__(self, config: DistilBertConfig, **kwargs):
|
||||
super().__init__(config, **kwargs)
|
||||
self.num_labels = config.num_labels
|
||||
self.config = config
|
||||
|
||||
self.distilbert = DistilBertModel(config)
|
||||
self.pre_classifier = nn.Linear(config.dim, config.dim)
|
||||
self.classifier = nn.Linear(config.dim, config.num_labels)
|
||||
self.dropout = nn.Dropout(config.seq_classif_dropout)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def get_position_embeddings(self) -> nn.Embedding:
|
||||
"""
|
||||
Returns the position embeddings
|
||||
"""
|
||||
return self.distilbert.get_position_embeddings()
|
||||
|
||||
def resize_position_embeddings(self, new_num_position_embeddings: int):
|
||||
"""
|
||||
Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.
|
||||
|
||||
Arguments:
|
||||
new_num_position_embeddings (`int`):
|
||||
The number of new position embedding matrix. If position embeddings are learned, increasing the size
|
||||
will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the
|
||||
end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the
|
||||
size will add correct vectors at the end following the position encoding algorithm, whereas reducing
|
||||
the size will remove vectors from the end.
|
||||
"""
|
||||
self.distilbert.resize_position_embeddings(new_num_position_embeddings)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
head_mask: Optional[torch.Tensor] = None,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
) -> Union[SequenceClassifierOutput, Tuple[torch.Tensor, ...]]:
|
||||
r"""
|
||||
labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
|
||||
Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
|
||||
config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
|
||||
`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
|
||||
"""
|
||||
return_dict = (
|
||||
return_dict if return_dict is not None else self.config.use_return_dict
|
||||
)
|
||||
|
||||
distilbert_output = self.distilbert(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
head_mask=head_mask,
|
||||
inputs_embeds=inputs_embeds,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
hidden_state = distilbert_output[0] # (bs, seq_len, dim)
|
||||
pooled_output = hidden_state[:, 0] # (bs, dim)
|
||||
pooled_output = self.pre_classifier(pooled_output) # (bs, dim)
|
||||
pooled_output = nn.ReLU()(pooled_output) # (bs, dim)
|
||||
pooled_output = self.dropout(pooled_output) # (bs, dim)
|
||||
logits = self.classifier(pooled_output) # (bs, num_labels)
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
if self.config.problem_type is None:
|
||||
if self.num_labels == 1:
|
||||
self.config.problem_type = "regression"
|
||||
elif self.num_labels > 1 and (
|
||||
labels.dtype == torch.long or labels.dtype == torch.int
|
||||
):
|
||||
self.config.problem_type = "single_label_classification"
|
||||
else:
|
||||
self.config.problem_type = "multi_label_classification"
|
||||
|
||||
if self.config.problem_type == "regression":
|
||||
loss_fct = MSELoss()
|
||||
if self.num_labels == 1:
|
||||
loss = loss_fct(logits.squeeze(), labels.squeeze())
|
||||
else:
|
||||
loss = loss_fct(logits, labels)
|
||||
elif self.config.problem_type == "single_label_classification":
|
||||
loss_fct = CrossEntropyLoss()
|
||||
loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
|
||||
elif self.config.problem_type == "multi_label_classification":
|
||||
loss_fct = BCEWithLogitsLoss()
|
||||
loss = loss_fct(logits, labels)
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + distilbert_output[1:]
|
||||
return ((loss,) + output) if loss is not None else output
|
||||
|
||||
return SequenceClassifierOutput(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
hidden_states=distilbert_output.hidden_states,
|
||||
attentions=distilbert_output.attentions,
|
||||
)
|
||||
@@ -0,0 +1,8 @@
|
||||
from typing import List
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class OCRErrorDetectionResult(BaseModel):
|
||||
texts: List[str]
|
||||
labels: List[str]
|
||||
@@ -0,0 +1,525 @@
|
||||
import collections
|
||||
import os
|
||||
import unicodedata
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from transformers.tokenization_utils import (
|
||||
PreTrainedTokenizer,
|
||||
_is_control,
|
||||
_is_punctuation,
|
||||
_is_whitespace,
|
||||
)
|
||||
|
||||
from surya.common.s3 import S3DownloaderMixin
|
||||
|
||||
VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
|
||||
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.load_vocab
|
||||
def load_vocab(vocab_file):
|
||||
"""Loads a vocabulary file into a dictionary."""
|
||||
vocab = collections.OrderedDict()
|
||||
with open(vocab_file, "r", encoding="utf-8") as reader:
|
||||
tokens = reader.readlines()
|
||||
for index, token in enumerate(tokens):
|
||||
token = token.rstrip("\n")
|
||||
vocab[token] = index
|
||||
return vocab
|
||||
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.whitespace_tokenize
|
||||
def whitespace_tokenize(text):
|
||||
"""Runs basic whitespace cleaning and splitting on a piece of text."""
|
||||
text = text.strip()
|
||||
if not text:
|
||||
return []
|
||||
tokens = text.split()
|
||||
return tokens
|
||||
|
||||
|
||||
class DistilBertTokenizer(S3DownloaderMixin, PreTrainedTokenizer):
|
||||
r"""
|
||||
Construct a DistilBERT tokenizer. Based on WordPiece.
|
||||
|
||||
This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
|
||||
this superclass for more information regarding those methods.
|
||||
|
||||
Args:
|
||||
vocab_file (`str`):
|
||||
File containing the vocabulary.
|
||||
do_lower_case (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to lowercase the input when tokenizing.
|
||||
do_basic_tokenize (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to do basic tokenization before WordPiece.
|
||||
never_split (`Iterable`, *optional*):
|
||||
Collection of tokens which will never be split during tokenization. Only has an effect when
|
||||
`do_basic_tokenize=True`
|
||||
unk_token (`str`, *optional*, defaults to `"[UNK]"`):
|
||||
The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
|
||||
token instead.
|
||||
sep_token (`str`, *optional*, defaults to `"[SEP]"`):
|
||||
The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
|
||||
sequence classification or for a text and a question for question answering. It is also used as the last
|
||||
token of a sequence built with special tokens.
|
||||
pad_token (`str`, *optional*, defaults to `"[PAD]"`):
|
||||
The token used for padding, for example when batching sequences of different lengths.
|
||||
cls_token (`str`, *optional*, defaults to `"[CLS]"`):
|
||||
The classifier token which is used when doing sequence classification (classification of the whole sequence
|
||||
instead of per-token classification). It is the first token of the sequence when built with special tokens.
|
||||
mask_token (`str`, *optional*, defaults to `"[MASK]"`):
|
||||
The token used for masking values. This is the token used when training this model with masked language
|
||||
modeling. This is the token which the model will try to predict.
|
||||
tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to tokenize Chinese characters.
|
||||
|
||||
This should likely be deactivated for Japanese (see this
|
||||
[issue](https://github.com/huggingface/transformers/issues/328)).
|
||||
strip_accents (`bool`, *optional*):
|
||||
Whether or not to strip all accents. If this option is not specified, then it will be determined by the
|
||||
value for `lowercase` (as in the original BERT).
|
||||
"""
|
||||
|
||||
vocab_files_names = VOCAB_FILES_NAMES
|
||||
model_input_names = ["input_ids", "attention_mask"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_file,
|
||||
do_lower_case=True,
|
||||
do_basic_tokenize=True,
|
||||
never_split=None,
|
||||
unk_token="[UNK]",
|
||||
sep_token="[SEP]",
|
||||
pad_token="[PAD]",
|
||||
cls_token="[CLS]",
|
||||
mask_token="[MASK]",
|
||||
tokenize_chinese_chars=True,
|
||||
strip_accents=None,
|
||||
**kwargs,
|
||||
):
|
||||
if not os.path.isfile(vocab_file):
|
||||
raise ValueError(
|
||||
f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
|
||||
" model use `tokenizer = DistilBertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
|
||||
)
|
||||
self.vocab = load_vocab(vocab_file)
|
||||
self.ids_to_tokens = collections.OrderedDict(
|
||||
[(ids, tok) for tok, ids in self.vocab.items()]
|
||||
)
|
||||
self.do_basic_tokenize = do_basic_tokenize
|
||||
if do_basic_tokenize:
|
||||
self.basic_tokenizer = BasicTokenizer(
|
||||
do_lower_case=do_lower_case,
|
||||
never_split=never_split,
|
||||
tokenize_chinese_chars=tokenize_chinese_chars,
|
||||
strip_accents=strip_accents,
|
||||
)
|
||||
self.wordpiece_tokenizer = WordpieceTokenizer(
|
||||
vocab=self.vocab, unk_token=str(unk_token)
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
do_lower_case=do_lower_case,
|
||||
do_basic_tokenize=do_basic_tokenize,
|
||||
never_split=never_split,
|
||||
unk_token=unk_token,
|
||||
sep_token=sep_token,
|
||||
pad_token=pad_token,
|
||||
cls_token=cls_token,
|
||||
mask_token=mask_token,
|
||||
tokenize_chinese_chars=tokenize_chinese_chars,
|
||||
strip_accents=strip_accents,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@property
|
||||
# Copied from transformers.models.bert.tokenization_bert.BertTokenizer.do_lower_case
|
||||
def do_lower_case(self):
|
||||
return self.basic_tokenizer.do_lower_case
|
||||
|
||||
@property
|
||||
# Copied from transformers.models.bert.tokenization_bert.BertTokenizer.vocab_size
|
||||
def vocab_size(self):
|
||||
return len(self.vocab)
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_vocab
|
||||
def get_vocab(self):
|
||||
return dict(self.vocab, **self.added_tokens_encoder)
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.BertTokenizer._tokenize
|
||||
def _tokenize(self, text, split_special_tokens=False):
|
||||
split_tokens = []
|
||||
if self.do_basic_tokenize:
|
||||
for token in self.basic_tokenizer.tokenize(
|
||||
text,
|
||||
never_split=self.all_special_tokens
|
||||
if not split_special_tokens
|
||||
else None,
|
||||
):
|
||||
# If the token is part of the never_split set
|
||||
if token in self.basic_tokenizer.never_split:
|
||||
split_tokens.append(token)
|
||||
else:
|
||||
split_tokens += self.wordpiece_tokenizer.tokenize(token)
|
||||
else:
|
||||
split_tokens = self.wordpiece_tokenizer.tokenize(text)
|
||||
return split_tokens
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_token_to_id
|
||||
def _convert_token_to_id(self, token):
|
||||
"""Converts a token (str) in an id using the vocab."""
|
||||
return self.vocab.get(token, self.vocab.get(self.unk_token))
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_id_to_token
|
||||
def _convert_id_to_token(self, index):
|
||||
"""Converts an index (integer) in a token (str) using the vocab."""
|
||||
return self.ids_to_tokens.get(index, self.unk_token)
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.BertTokenizer.convert_tokens_to_string
|
||||
def convert_tokens_to_string(self, tokens):
|
||||
"""Converts a sequence of tokens (string) in a single string."""
|
||||
out_string = " ".join(tokens).replace(" ##", "").strip()
|
||||
return out_string
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.BertTokenizer.build_inputs_with_special_tokens
|
||||
def build_inputs_with_special_tokens(
|
||||
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
|
||||
) -> List[int]:
|
||||
"""
|
||||
Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
|
||||
adding special tokens. A BERT sequence has the following format:
|
||||
|
||||
- single sequence: `[CLS] X [SEP]`
|
||||
- pair of sequences: `[CLS] A [SEP] B [SEP]`
|
||||
|
||||
Args:
|
||||
token_ids_0 (`List[int]`):
|
||||
List of IDs to which the special tokens will be added.
|
||||
token_ids_1 (`List[int]`, *optional*):
|
||||
Optional second list of IDs for sequence pairs.
|
||||
|
||||
Returns:
|
||||
`List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
|
||||
"""
|
||||
if token_ids_1 is None:
|
||||
return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
|
||||
cls = [self.cls_token_id]
|
||||
sep = [self.sep_token_id]
|
||||
return cls + token_ids_0 + sep + token_ids_1 + sep
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_special_tokens_mask
|
||||
def get_special_tokens_mask(
|
||||
self,
|
||||
token_ids_0: List[int],
|
||||
token_ids_1: Optional[List[int]] = None,
|
||||
already_has_special_tokens: bool = False,
|
||||
) -> List[int]:
|
||||
"""
|
||||
Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
|
||||
special tokens using the tokenizer `prepare_for_model` method.
|
||||
|
||||
Args:
|
||||
token_ids_0 (`List[int]`):
|
||||
List of IDs.
|
||||
token_ids_1 (`List[int]`, *optional*):
|
||||
Optional second list of IDs for sequence pairs.
|
||||
already_has_special_tokens (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not the token list is already formatted with special tokens for the model.
|
||||
|
||||
Returns:
|
||||
`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
|
||||
"""
|
||||
|
||||
if already_has_special_tokens:
|
||||
return super().get_special_tokens_mask(
|
||||
token_ids_0=token_ids_0,
|
||||
token_ids_1=token_ids_1,
|
||||
already_has_special_tokens=True,
|
||||
)
|
||||
|
||||
if token_ids_1 is not None:
|
||||
return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
|
||||
return [1] + ([0] * len(token_ids_0)) + [1]
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.BertTokenizer.create_token_type_ids_from_sequences
|
||||
def create_token_type_ids_from_sequences(
|
||||
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
|
||||
) -> List[int]:
|
||||
"""
|
||||
Create a mask from the two sequences passed to be used in a sequence-pair classification task. A BERT sequence
|
||||
pair mask has the following format:
|
||||
|
||||
```
|
||||
0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
|
||||
| first sequence | second sequence |
|
||||
```
|
||||
|
||||
If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
|
||||
|
||||
Args:
|
||||
token_ids_0 (`List[int]`):
|
||||
List of IDs.
|
||||
token_ids_1 (`List[int]`, *optional*):
|
||||
Optional second list of IDs for sequence pairs.
|
||||
|
||||
Returns:
|
||||
`List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
|
||||
"""
|
||||
sep = [self.sep_token_id]
|
||||
cls = [self.cls_token_id]
|
||||
if token_ids_1 is None:
|
||||
return len(cls + token_ids_0 + sep) * [0]
|
||||
return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.BertTokenizer.save_vocabulary
|
||||
def save_vocabulary(
|
||||
self, save_directory: str, filename_prefix: Optional[str] = None
|
||||
) -> Tuple[str]:
|
||||
index = 0
|
||||
if os.path.isdir(save_directory):
|
||||
vocab_file = os.path.join(
|
||||
save_directory,
|
||||
(filename_prefix + "-" if filename_prefix else "")
|
||||
+ VOCAB_FILES_NAMES["vocab_file"],
|
||||
)
|
||||
else:
|
||||
vocab_file = (
|
||||
filename_prefix + "-" if filename_prefix else ""
|
||||
) + save_directory
|
||||
with open(vocab_file, "w", encoding="utf-8") as writer:
|
||||
for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
|
||||
if index != token_index:
|
||||
# logger.warning(
|
||||
# f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
|
||||
# " Please check that the vocabulary is not corrupted!"
|
||||
# )
|
||||
index = token_index
|
||||
writer.write(token + "\n")
|
||||
index += 1
|
||||
return (vocab_file,)
|
||||
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.BasicTokenizer
|
||||
class BasicTokenizer(object):
|
||||
"""
|
||||
Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).
|
||||
|
||||
Args:
|
||||
do_lower_case (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to lowercase the input when tokenizing.
|
||||
never_split (`Iterable`, *optional*):
|
||||
Collection of tokens which will never be split during tokenization. Only has an effect when
|
||||
`do_basic_tokenize=True`
|
||||
tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to tokenize Chinese characters.
|
||||
|
||||
This should likely be deactivated for Japanese (see this
|
||||
[issue](https://github.com/huggingface/transformers/issues/328)).
|
||||
strip_accents (`bool`, *optional*):
|
||||
Whether or not to strip all accents. If this option is not specified, then it will be determined by the
|
||||
value for `lowercase` (as in the original BERT).
|
||||
do_split_on_punc (`bool`, *optional*, defaults to `True`):
|
||||
In some instances we want to skip the basic punctuation splitting so that later tokenization can capture
|
||||
the full context of the words, such as contractions.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
do_lower_case=True,
|
||||
never_split=None,
|
||||
tokenize_chinese_chars=True,
|
||||
strip_accents=None,
|
||||
do_split_on_punc=True,
|
||||
):
|
||||
if never_split is None:
|
||||
never_split = []
|
||||
self.do_lower_case = do_lower_case
|
||||
self.never_split = set(never_split)
|
||||
self.tokenize_chinese_chars = tokenize_chinese_chars
|
||||
self.strip_accents = strip_accents
|
||||
self.do_split_on_punc = do_split_on_punc
|
||||
|
||||
def tokenize(self, text, never_split=None):
|
||||
"""
|
||||
Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.
|
||||
|
||||
Args:
|
||||
never_split (`List[str]`, *optional*)
|
||||
Kept for backward compatibility purposes. Now implemented directly at the base class level (see
|
||||
[`PreTrainedTokenizer.tokenize`]) List of token not to split.
|
||||
"""
|
||||
# union() returns a new set by concatenating the two sets.
|
||||
never_split = (
|
||||
self.never_split.union(set(never_split))
|
||||
if never_split
|
||||
else self.never_split
|
||||
)
|
||||
text = self._clean_text(text)
|
||||
|
||||
# This was added on November 1st, 2018 for the multilingual and Chinese
|
||||
# models. This is also applied to the English models now, but it doesn't
|
||||
# matter since the English models were not trained on any Chinese data
|
||||
# and generally don't have any Chinese data in them (there are Chinese
|
||||
# characters in the vocabulary because Wikipedia does have some Chinese
|
||||
# words in the English Wikipedia.).
|
||||
if self.tokenize_chinese_chars:
|
||||
text = self._tokenize_chinese_chars(text)
|
||||
# prevents treating the same character with different unicode codepoints as different characters
|
||||
unicode_normalized_text = unicodedata.normalize("NFC", text)
|
||||
orig_tokens = whitespace_tokenize(unicode_normalized_text)
|
||||
split_tokens = []
|
||||
for token in orig_tokens:
|
||||
if token not in never_split:
|
||||
if self.do_lower_case:
|
||||
token = token.lower()
|
||||
if self.strip_accents is not False:
|
||||
token = self._run_strip_accents(token)
|
||||
elif self.strip_accents:
|
||||
token = self._run_strip_accents(token)
|
||||
split_tokens.extend(self._run_split_on_punc(token, never_split))
|
||||
|
||||
output_tokens = whitespace_tokenize(" ".join(split_tokens))
|
||||
return output_tokens
|
||||
|
||||
def _run_strip_accents(self, text):
|
||||
"""Strips accents from a piece of text."""
|
||||
text = unicodedata.normalize("NFD", text)
|
||||
output = []
|
||||
for char in text:
|
||||
cat = unicodedata.category(char)
|
||||
if cat == "Mn":
|
||||
continue
|
||||
output.append(char)
|
||||
return "".join(output)
|
||||
|
||||
def _run_split_on_punc(self, text, never_split=None):
|
||||
"""Splits punctuation on a piece of text."""
|
||||
if not self.do_split_on_punc or (
|
||||
never_split is not None and text in never_split
|
||||
):
|
||||
return [text]
|
||||
chars = list(text)
|
||||
i = 0
|
||||
start_new_word = True
|
||||
output = []
|
||||
while i < len(chars):
|
||||
char = chars[i]
|
||||
if _is_punctuation(char):
|
||||
output.append([char])
|
||||
start_new_word = True
|
||||
else:
|
||||
if start_new_word:
|
||||
output.append([])
|
||||
start_new_word = False
|
||||
output[-1].append(char)
|
||||
i += 1
|
||||
|
||||
return ["".join(x) for x in output]
|
||||
|
||||
def _tokenize_chinese_chars(self, text):
|
||||
"""Adds whitespace around any CJK character."""
|
||||
output = []
|
||||
for char in text:
|
||||
cp = ord(char)
|
||||
if self._is_chinese_char(cp):
|
||||
output.append(" ")
|
||||
output.append(char)
|
||||
output.append(" ")
|
||||
else:
|
||||
output.append(char)
|
||||
return "".join(output)
|
||||
|
||||
def _is_chinese_char(self, cp):
|
||||
"""Checks whether CP is the codepoint of a CJK character."""
|
||||
# This defines a "chinese character" as anything in the CJK Unicode block:
|
||||
# https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
|
||||
#
|
||||
# Note that the CJK Unicode block is NOT all Japanese and Korean characters,
|
||||
# despite its name. The modern Korean Hangul alphabet is a different block,
|
||||
# as is Japanese Hiragana and Katakana. Those alphabets are used to write
|
||||
# space-separated words, so they are not treated specially and handled
|
||||
# like the all of the other languages.
|
||||
if (
|
||||
(cp >= 0x4E00 and cp <= 0x9FFF)
|
||||
or (cp >= 0x3400 and cp <= 0x4DBF) #
|
||||
or (cp >= 0x20000 and cp <= 0x2A6DF) #
|
||||
or (cp >= 0x2A700 and cp <= 0x2B73F) #
|
||||
or (cp >= 0x2B740 and cp <= 0x2B81F) #
|
||||
or (cp >= 0x2B820 and cp <= 0x2CEAF) #
|
||||
or (cp >= 0xF900 and cp <= 0xFAFF)
|
||||
or (cp >= 0x2F800 and cp <= 0x2FA1F) #
|
||||
): #
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _clean_text(self, text):
|
||||
"""Performs invalid character removal and whitespace cleanup on text."""
|
||||
output = []
|
||||
for char in text:
|
||||
cp = ord(char)
|
||||
if cp == 0 or cp == 0xFFFD or _is_control(char):
|
||||
continue
|
||||
if _is_whitespace(char):
|
||||
output.append(" ")
|
||||
else:
|
||||
output.append(char)
|
||||
return "".join(output)
|
||||
|
||||
|
||||
# Copied from transformers.models.bert.tokenization_bert.WordpieceTokenizer
|
||||
class WordpieceTokenizer(object):
|
||||
"""Runs WordPiece tokenization."""
|
||||
|
||||
def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
|
||||
self.vocab = vocab
|
||||
self.unk_token = unk_token
|
||||
self.max_input_chars_per_word = max_input_chars_per_word
|
||||
|
||||
def tokenize(self, text):
|
||||
"""
|
||||
Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
|
||||
tokenization using the given vocabulary.
|
||||
|
||||
For example, `input = "unaffable"` wil return as output `["un", "##aff", "##able"]`.
|
||||
|
||||
Args:
|
||||
text: A single token or whitespace separated tokens. This should have
|
||||
already been passed through *BasicTokenizer*.
|
||||
|
||||
Returns:
|
||||
A list of wordpiece tokens.
|
||||
"""
|
||||
|
||||
output_tokens = []
|
||||
for token in whitespace_tokenize(text):
|
||||
chars = list(token)
|
||||
if len(chars) > self.max_input_chars_per_word:
|
||||
output_tokens.append(self.unk_token)
|
||||
continue
|
||||
|
||||
is_bad = False
|
||||
start = 0
|
||||
sub_tokens = []
|
||||
while start < len(chars):
|
||||
end = len(chars)
|
||||
cur_substr = None
|
||||
while start < end:
|
||||
substr = "".join(chars[start:end])
|
||||
if start > 0:
|
||||
substr = "##" + substr
|
||||
if substr in self.vocab:
|
||||
cur_substr = substr
|
||||
break
|
||||
end -= 1
|
||||
if cur_substr is None:
|
||||
is_bad = True
|
||||
break
|
||||
sub_tokens.append(cur_substr)
|
||||
start = end
|
||||
|
||||
if is_bad:
|
||||
output_tokens.append(self.unk_token)
|
||||
else:
|
||||
output_tokens.extend(sub_tokens)
|
||||
return output_tokens
|
||||
Reference in New Issue
Block a user