import argparse
import json
from pathlib import Path
import re
from typing import Dict, Optional, Union
import torch
import torch.nn.functional as F
from .modules.layers import LstmSeq2SeqEncoder
from .modules.base import InstructBase
from .modules.evaluator import Evaluator, greedy_search
from .modules.span_rep import SpanRepLayer
from .modules.token_rep import TokenRepLayer
from torch import nn
from torch.nn.utils.rnn import pad_sequence
from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
from huggingface_hub.utils import HfHubHTTPError
class GLiNER(InstructBase, PyTorchModelHubMixin):
def __init__(self, config):
super().__init__(config)
self.config = config
# [ENT] token
self.entity_token = "<<ENT>>"
self.sep_token = "<<SEP>>"
# usually a pretrained bidirectional transformer, returns first subtoken representation
self.token_rep_layer = TokenRepLayer(model_name=config.model_name, fine_tune=config.fine_tune,
subtoken_pooling=config.subtoken_pooling, hidden_size=config.hidden_size,
add_tokens=[self.entity_token, self.sep_token])
# hierarchical representation of tokens
self.rnn = LstmSeq2SeqEncoder(
input_size=config.hidden_size,
hidden_size=config.hidden_size // 2,
num_layers=1,
bidirectional=True,
)
# span representation
self.span_rep_layer = SpanRepLayer(
span_mode=config.span_mode,
hidden_size=config.hidden_size,
max_width=config.max_width,
dropout=config.dropout,
)
# prompt representation (FFN)
self.prompt_rep_layer = nn.Sequential(
nn.Linear(config.hidden_size, config.hidden_size * 4),
nn.Dropout(config.dropout),
nn.ReLU(),
nn.Linear(config.hidden_size * 4, config.hidden_size)
)
def compute_score_train(self, x):
span_idx = x['span_idx'] * x['span_mask'].unsqueeze(-1)
new_length = x['seq_length'].clone()
new_tokens = []
all_len_prompt = []
num_classes_all = []
# add prompt to the tokens
for i in range(len(x['tokens'])):
all_types_i = list(x['classes_to_id'][i].keys())
# multiple entity types in all_types. Prompt is appended at the start of tokens
entity_prompt = []
num_classes_all.append(len(all_types_i))
# add enity types to prompt
for entity_type in all_types_i:
entity_prompt.append(self.entity_token) # [ENT] token
entity_prompt.append(entity_type) # entity type
entity_prompt.append(self.sep_token) # [SEP] token
# prompt format:
# [ENT] entity_type [ENT] entity_type ... [ENT] entity_type [SEP]
# add prompt to the tokens
tokens_p = entity_prompt + x['tokens'][i]
# input format:
# [ENT] entity_type_1 [ENT] entity_type_2 ... [ENT] entity_type_m [SEP] token_1 token_2 ... token_n
# update length of the sequence (add prompt length to the original length)
new_length[i] = new_length[i] + len(entity_prompt)
# update tokens
new_tokens.append(tokens_p)
# store prompt length
all_len_prompt.append(len(entity_prompt))
# create a mask using num_classes_all (0, if it exceeds the number of classes, 1 otherwise)
max_num_classes = max(num_classes_all)
entity_type_mask = torch.arange(max_num_classes).unsqueeze(0).expand(len(num_classes_all), -1).to(
x['span_mask'].device)
entity_type_mask = entity_type_mask < torch.tensor(num_classes_all).unsqueeze(-1).to(
x['span_mask'].device) # [batch_size, max_num_classes]
# compute all token representations
bert_output = self.token_rep_layer(new_tokens, new_length)
word_rep_w_prompt = bert_output["embeddings"] # embeddings for all tokens (with prompt)
mask_w_prompt = bert_output["mask"] # mask for all tokens (with prompt)
# get word representation (after [SEP]), mask (after [SEP]) and entity type representation (before [SEP])
word_rep = [] # word representation (after [SEP])
mask = [] # mask (after [SEP])
entity_type_rep = [] # entity type representation (before [SEP])
for i in range(len(x['tokens'])):
prompt_entity_length = all_len_prompt[i] # length of prompt for this example
# get word representation (after [SEP])
word_rep.append(word_rep_w_prompt[i, prompt_entity_length:prompt_entity_length + x['seq_length'][i]])
# get mask (after [SEP])
mask.append(mask_w_prompt[i, prompt_entity_length:prompt_entity_length + x['seq_length'][i]])
# get entity type representation (before [SEP])
entity_rep = word_rep_w_prompt[i, :prompt_entity_length - 1] # remove [SEP]
entity_rep = entity_rep[0::2] # it means that we take every second element starting from the second one
entity_type_rep.append(entity_rep)
# padding for word_rep, mask and entity_type_rep
word_rep = pad_sequence(word_rep, batch_first=True) # [batch_size, seq_len, hidden_size]
mask = pad_sequence(mask, batch_first=True) # [batch_size, seq_len]
entity_type_rep = pad_sequence(entity_type_rep, batch_first=True) # [batch_size, len_types, hidden_size]
# compute span representation
word_rep = self.rnn(word_rep, mask)
span_rep = self.span_rep_layer(word_rep, span_idx)
# compute final entity type representation (FFN)
entity_type_rep = self.prompt_rep_layer(entity_type_rep) # (batch_size, len_types, hidden_size)
num_classes = entity_type_rep.shape[1] # number of entity types
# similarity score
scores = torch.einsum('BLKD,BCD->BLKC', span_rep, entity_type_rep)
return scores, num_classes, entity_type_mask
def forward(self, x):
# compute span representation
scores, num_classes, entity_type_mask = self.compute_score_train(x)
batch_size = scores.shape[0]
# loss for filtering classifier
logits_label = scores.view(-1, num_classes)
labels = x["span_label"].view(-1) # (batch_size * num_spans)
mask_label = labels != -1 # (batch_size * num_spans)
labels.masked_fill_(~mask_label, 0) # Set the labels of padding tokens to 0
# one-hot encoding
labels_one_hot = torch.zeros(labels.size(0), num_classes + 1, dtype=torch.float32).to(scores.device)
labels_one_hot.scatter_(1, labels.unsqueeze(1), 1) # Set the corresponding index to 1
labels_one_hot = labels_one_hot[:, 1:] # Remove the first column
# Shape of labels_one_hot: (batch_size * num_spans, num_classes)
# compute loss (without reduction)
all_losses = F.binary_cross_entropy_with_logits(logits_label, labels_one_hot,
reduction='none')
# mask loss using entity_type_mask (B, C)
masked_loss = all_losses.view(batch_size, -1, num_classes) * entity_type_mask.unsqueeze(1)
all_losses = masked_loss.view(-1, num_classes)
# expand mask_label to all_losses
mask_label = mask_label.unsqueeze(-1).expand_as(all_losses)
# put lower loss for in label_one_hot (2 for positive, 1 for negative)
weight_c = labels_one_hot + 1
# apply mask
all_losses = all_losses * mask_label.float() * weight_c
return all_losses.sum()
def compute_score_eval(self, x, device):
# check if classes_to_id is dict
assert isinstance(x['classes_to_id'], dict), "classes_to_id must be a dict"
span_idx = (x['span_idx'] * x['span_mask'].unsqueeze(-1)).to(device)
all_types = list(x['classes_to_id'].keys())
# multiple entity types in all_types. Prompt is appended at the start of tokens
entity_prompt = []
# add enity types to prompt
for entity_type in all_types:
entity_prompt.append(self.entity_token)
entity_prompt.append(entity_type)
entity_prompt.append(self.sep_token)
prompt_entity_length = len(entity_prompt)
# add prompt
tokens_p = [entity_prompt + tokens for tokens in x['tokens']]
seq_length_p = x['seq_length'] + prompt_entity_length
out = self.token_rep_layer(tokens_p, seq_length_p)
word_rep_w_prompt = out["embeddings"]
mask_w_prompt = out["mask"]
# remove prompt
word_rep = word_rep_w_prompt[:, prompt_entity_length:, :]
mask = mask_w_prompt[:, prompt_entity_length:]
# get_entity_type_rep
entity_type_rep = word_rep_w_prompt[:, :prompt_entity_length - 1, :]
# extract [ENT] tokens (which are at even positions in entity_type_rep)
entity_type_rep = entity_type_rep[:, 0::2, :]
entity_type_rep = self.prompt_rep_layer(entity_type_rep) # (batch_size, len_types, hidden_size)
word_rep = self.rnn(word_rep, mask)
span_rep = self.span_rep_layer(word_rep, span_idx)
local_scores = torch.einsum('BLKD,BCD->BLKC', span_rep, entity_type_rep)
return local_scores
@torch.no_grad()
def predict(self, x, flat_ner=False, threshold=0.5):
self.eval()
local_scores = self.compute_score_eval(x, device=next(self.parameters()).device)
spans = []
for i, _ in enumerate(x["tokens"]):
local_i = local_scores[i]
wh_i = [i.tolist() for i in torch.where(torch.sigmoid(local_i) > threshold)]
span_i = []
for s, k, c in zip(*wh_i):
if s + k < len(x["tokens"][i]):
span_i.append((s, s + k, x["id_to_classes"][c + 1], local_i[s, k, c]))
span_i = greedy_search(span_i, flat_ner)
spans.append(span_i)
return spans
def predict_entities(self, text, labels, flat_ner=True, threshold=0.5):
tokens = []
start_token_idx_to_text_idx = []
end_token_idx_to_text_idx = []
for match in re.finditer(r'\w+(?:[-_]\w+)*|\S', text):
tokens.append(match.group())
start_token_idx_to_text_idx.append(match.start())
end_token_idx_to_text_idx.append(match.end())
input_x = {"tokenized_text": tokens, "ner": None}
x = self.collate_fn([input_x], labels)
output = self.predict(x, flat_ner=flat_ner, threshold=threshold)
entities = []
for start_token_idx, end_token_idx, ent_type in output[0]:
start_text_idx = start_token_idx_to_text_idx[start_token_idx]
end_text_idx = end_token_idx_to_text_idx[end_token_idx]
entities.append({
"start": start_token_idx_to_text_idx[start_token_idx],
"end": end_token_idx_to_text_idx[end_token_idx],
"text": text[start_text_idx:end_text_idx],
"label": ent_type,
})
return entities
def evaluate(self, test_data, flat_ner=False, threshold=0.5, batch_size=12, entity_types=None):
self.eval()
data_loader = self.create_dataloader(test_data, batch_size=batch_size, entity_types=entity_types, shuffle=False)
device = next(self.parameters()).device
all_preds = []
all_trues = []
for x in data_loader:
for k, v in x.items():
if isinstance(v, torch.Tensor):
x[k] = v.to(device)
batch_predictions = self.predict(x, flat_ner, threshold)
all_preds.extend(batch_predictions)
all_trues.extend(x["entities"])
evaluator = Evaluator(all_trues, all_preds)
out, f1 = evaluator.evaluate()
return out, f1
@classmethod
def _from_pretrained(
cls,
*,
model_id: str,
revision: Optional[str],
cache_dir: Optional[Union[str, Path]],
force_download: bool,
proxies: Optional[Dict],
resume_download: bool,
local_files_only: bool,
token: Union[str, bool, None],
map_location: str = "cpu",
strict: bool = False,
**model_kwargs,
):
# 1. Backwards compatibility: Use "gliner_base.pt" and "gliner_multi.pt" with all data
filenames = ["gliner_base.pt", "gliner_multi.pt"]
for filename in filenames:
model_file = Path(model_id) / filename
if not model_file.exists():
try:
model_file = hf_hub_download(
repo_id=model_id,
filename=filename,
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
token=token,
local_files_only=local_files_only,
)
except HfHubHTTPError:
continue
dict_load = torch.load(model_file, map_location=torch.device(map_location))
config = dict_load["config"]
state_dict = dict_load["model_weights"]
config.model_name = "microsoft/deberta-v3-base" if filename == "gliner_base.pt" else "microsoft/mdeberta-v3-base"
model = cls(config)
model.load_state_dict(state_dict, strict=strict, assign=True)
# Required to update flair's internals as well:
model.to(map_location)
return model
# 2. Newer format: Use "pytorch_model.bin" and "gliner_config.json"
from .train import load_config_as_namespace
model_file = Path(model_id) / "pytorch_model.bin"
if not model_file.exists():
model_file = hf_hub_download(
repo_id=model_id,
filename="pytorch_model.bin",
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
token=token,
local_files_only=local_files_only,
)
config_file = Path(model_id) / "gliner_config.json"
if not config_file.exists():
config_file = hf_hub_download(
repo_id=model_id,
filename="gliner_config.json",
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
token=token,
local_files_only=local_files_only,
)
config = load_config_as_namespace(config_file)
model = cls(config)
state_dict = torch.load(model_file, map_location=torch.device(map_location))
model.load_state_dict(state_dict, strict=strict, assign=True)
model.to(map_location)
return model
def save_pretrained(
self,
save_directory: Union[str, Path],
*,
config: Optional[Union[dict, "DataclassInstance"]] = None,
repo_id: Optional[str] = None,
push_to_hub: bool = False,
**push_to_hub_kwargs,
) -> Optional[str]:
"""
Save weights in local directory.
Args:
save_directory (`str` or `Path`):
Path to directory in which the model weights and configuration will be saved.
config (`dict` or `DataclassInstance`, *optional*):
Model configuration specified as a key/value dictionary or a dataclass instance.
push_to_hub (`bool`, *optional*, defaults to `False`):
Whether or not to push your model to the Huggingface Hub after saving it.
repo_id (`str`, *optional*):
ID of your repository on the Hub. Used only if `push_to_hub=True`. Will default to the folder name if
not provided.
kwargs:
Additional key word arguments passed along to the [`~ModelHubMixin.push_to_hub`] method.
"""
save_directory = Path(save_directory)
save_directory.mkdir(parents=True, exist_ok=True)
# save model weights/files
torch.save(self.state_dict(), save_directory / "pytorch_model.bin")
# save config (if provided)
if config is None:
config = self.config
if config is not None:
if isinstance(config, argparse.Namespace):
config = vars(config)
(save_directory / "gliner_config.json").write_text(json.dumps(config, indent=2))
# push to the Hub if required
if push_to_hub:
kwargs = push_to_hub_kwargs.copy() # soft-copy to avoid mutating input
if config is not None: # kwarg for `push_to_hub`
kwargs["config"] = config
if repo_id is None:
repo_id = save_directory.name # Defaults to `save_directory` name
return self.push_to_hub(repo_id=repo_id, **kwargs)
return None
def to(self, device):
super().to(device)
import flair
flair.device = device
return self