初始化项目，由ModelHub XC社区提供模型
Model: ByteDance/ListConRanker Source: Original Platform
This commit is contained in:
ModelHub XC
530
modeling_listconranker.py
Normal file
530
modeling_listconranker.py
Normal file
@@ -0,0 +1,530 @@
|
||||
# Copyright 2024 Bytedance Ltd. and/or its affiliates
|
||||
#
|
||||
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software
|
||||
# and associated documentation files (the "Software"), to deal in the Software without
|
||||
# restriction, including without limitation the rights to use, copy, modify, merge, publish,
|
||||
# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
|
||||
# Software is furnished to do so, subject to the following conditions:
|
||||
#
|
||||
# The above copyright notice and this permission notice shall be included in all copies or
|
||||
# substantial portions of the Software.
|
||||
#
|
||||
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
|
||||
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
|
||||
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
|
||||
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
|
||||
# OTHER DEALINGS IN THE SOFTWARE.
|
||||
from __future__ import annotations
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
from transformers import (
|
||||
PreTrainedModel,
|
||||
BertModel,
|
||||
AutoTokenizer,
|
||||
)
|
||||
import os
|
||||
from transformers.modeling_outputs import SequenceClassifierOutput
|
||||
from typing import Union, List, Optional
|
||||
from collections import defaultdict
|
||||
import numpy as np
|
||||
import math
|
||||
from huggingface_hub import hf_hub_download
|
||||
from .configuration_listconranker import ListConRankerConfig
|
||||
|
||||
|
||||
class QueryEmbedding(nn.Module):
|
||||
def __init__(self, config) -> None:
|
||||
super().__init__()
|
||||
self.query_embedding = nn.Embedding(2, config.list_con_hidden_size)
|
||||
self.layerNorm = nn.LayerNorm(config.list_con_hidden_size)
|
||||
|
||||
def forward(self, x, tags):
|
||||
query_embeddings = self.query_embedding(tags)
|
||||
x += query_embeddings
|
||||
x = self.layerNorm(x)
|
||||
return x
|
||||
|
||||
|
||||
class ListTransformer(nn.Module):
|
||||
def __init__(self, num_layer, config) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.list_transformer_layer = nn.TransformerEncoderLayer(
|
||||
config.list_con_hidden_size,
|
||||
self.config.num_attention_heads,
|
||||
batch_first=True,
|
||||
activation=F.gelu,
|
||||
norm_first=False,
|
||||
)
|
||||
self.list_transformer = nn.TransformerEncoder(
|
||||
self.list_transformer_layer, num_layer
|
||||
)
|
||||
self.relu = nn.ReLU()
|
||||
self.query_embedding = QueryEmbedding(config)
|
||||
|
||||
self.linear_score3 = nn.Linear(
|
||||
config.list_con_hidden_size * 2, config.list_con_hidden_size
|
||||
)
|
||||
self.linear_score2 = nn.Linear(
|
||||
config.list_con_hidden_size * 2, config.list_con_hidden_size
|
||||
)
|
||||
self.linear_score1 = nn.Linear(config.list_con_hidden_size * 2, 1)
|
||||
|
||||
def forward(
|
||||
self, pair_features: torch.Tensor, pair_nums: List[int]
|
||||
) -> torch.Tensor:
|
||||
batch_pair_features = pair_features.split(pair_nums)
|
||||
|
||||
pair_feature_query_passage_concat_list = []
|
||||
for i in range(len(batch_pair_features)):
|
||||
pair_feature_query = (
|
||||
batch_pair_features[i][0].unsqueeze(0).repeat(pair_nums[i] - 1, 1)
|
||||
)
|
||||
pair_feature_passage = batch_pair_features[i][1:]
|
||||
pair_feature_query_passage_concat_list.append(
|
||||
torch.cat([pair_feature_query, pair_feature_passage], dim=1)
|
||||
)
|
||||
pair_feature_query_passage_concat = torch.cat(
|
||||
pair_feature_query_passage_concat_list, dim=0
|
||||
)
|
||||
|
||||
batch_pair_features = nn.utils.rnn.pad_sequence(
|
||||
batch_pair_features, batch_first=True
|
||||
)
|
||||
|
||||
query_embedding_tags = torch.zeros(
|
||||
batch_pair_features.size(0),
|
||||
batch_pair_features.size(1),
|
||||
dtype=torch.long,
|
||||
device=self.device,
|
||||
)
|
||||
query_embedding_tags[:, 0] = 1
|
||||
batch_pair_features = self.query_embedding(
|
||||
batch_pair_features, query_embedding_tags
|
||||
)
|
||||
|
||||
mask = self.generate_attention_mask(pair_nums)
|
||||
query_mask = self.generate_attention_mask_custom(pair_nums)
|
||||
pair_list_features = self.list_transformer(
|
||||
batch_pair_features, src_key_padding_mask=mask, mask=query_mask
|
||||
)
|
||||
|
||||
output_pair_list_features = []
|
||||
output_query_list_features = []
|
||||
pair_features_after_transformer_list = []
|
||||
for idx, pair_num in enumerate(pair_nums):
|
||||
output_pair_list_features.append(pair_list_features[idx, 1:pair_num, :])
|
||||
output_query_list_features.append(pair_list_features[idx, 0, :])
|
||||
pair_features_after_transformer_list.append(
|
||||
pair_list_features[idx, :pair_num, :]
|
||||
)
|
||||
|
||||
pair_features_after_transformer_cat_query_list = []
|
||||
for idx, pair_num in enumerate(pair_nums):
|
||||
query_ft = (
|
||||
output_query_list_features[idx].unsqueeze(0).repeat(pair_num - 1, 1)
|
||||
)
|
||||
pair_features_after_transformer_cat_query = torch.cat(
|
||||
[query_ft, output_pair_list_features[idx]], dim=1
|
||||
)
|
||||
pair_features_after_transformer_cat_query_list.append(
|
||||
pair_features_after_transformer_cat_query
|
||||
)
|
||||
pair_features_after_transformer_cat_query = torch.cat(
|
||||
pair_features_after_transformer_cat_query_list, dim=0
|
||||
)
|
||||
|
||||
pair_feature_query_passage_concat = self.relu(
|
||||
self.linear_score2(pair_feature_query_passage_concat)
|
||||
)
|
||||
pair_features_after_transformer_cat_query = self.relu(
|
||||
self.linear_score3(pair_features_after_transformer_cat_query)
|
||||
)
|
||||
final_ft = torch.cat(
|
||||
[
|
||||
pair_feature_query_passage_concat,
|
||||
pair_features_after_transformer_cat_query,
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
logits = self.linear_score1(final_ft).squeeze()
|
||||
return logits, torch.cat(pair_features_after_transformer_list, dim=0)
|
||||
|
||||
def generate_attention_mask(self, pair_num):
|
||||
max_len = max(pair_num)
|
||||
batch_size = len(pair_num)
|
||||
mask = torch.zeros(batch_size, max_len, dtype=torch.bool, device=self.device)
|
||||
for i, length in enumerate(pair_num):
|
||||
mask[i, length:] = True
|
||||
return mask
|
||||
|
||||
def generate_attention_mask_custom(self, pair_num):
|
||||
max_len = max(pair_num)
|
||||
mask = torch.zeros(max_len, max_len, dtype=torch.bool, device=self.device)
|
||||
mask[0, 1:] = True
|
||||
return mask
|
||||
|
||||
|
||||
class ListConRankerModel(PreTrainedModel):
|
||||
"""
|
||||
ListConRanker model for sequence classification that's compatible with AutoModelForSequenceClassification.
|
||||
"""
|
||||
|
||||
config_class = ListConRankerConfig
|
||||
base_model_prefix = "listconranker"
|
||||
|
||||
def __init__(self, config: ListConRankerConfig):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
self.num_labels = config.num_labels
|
||||
self.hf_model = BertModel(config.bert_config)
|
||||
|
||||
self.sigmoid = nn.Sigmoid()
|
||||
|
||||
self.linear_in_embedding = nn.Linear(
|
||||
config.hidden_size, config.list_con_hidden_size
|
||||
)
|
||||
self.list_transformer = ListTransformer(
|
||||
config.list_transformer_layers,
|
||||
config,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
token_type_ids: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.Tensor] = None,
|
||||
head_mask: Optional[torch.Tensor] = None,
|
||||
inputs_embeds: Optional[torch.Tensor] = None,
|
||||
labels: Optional[torch.Tensor] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[tuple[torch.Tensor], SequenceClassifierOutput]:
|
||||
if self.training:
|
||||
raise NotImplementedError("Training not supported; use eval mode.")
|
||||
device = input_ids.device
|
||||
self.list_transformer.device = device
|
||||
# Reorganize by unique queries and their passages
|
||||
(
|
||||
reorganized_input_ids,
|
||||
reorganized_attention_mask,
|
||||
reorganized_token_type_ids,
|
||||
pair_nums,
|
||||
group_indices,
|
||||
) = self._reorganize_inputs(input_ids, attention_mask, token_type_ids)
|
||||
|
||||
out = self.hf_model(
|
||||
input_ids=reorganized_input_ids,
|
||||
attention_mask=reorganized_attention_mask,
|
||||
token_type_ids=reorganized_token_type_ids,
|
||||
return_dict=True,
|
||||
)
|
||||
feats = out.last_hidden_state
|
||||
pooled = self.average_pooling(feats, reorganized_attention_mask)
|
||||
embedded = self.linear_in_embedding(pooled)
|
||||
logits, _ = self.list_transformer(embedded, pair_nums)
|
||||
probs = self.sigmoid(logits)
|
||||
|
||||
# Restore original order
|
||||
sorted_probs = self._restore_original_order(probs, group_indices)
|
||||
sorted_logits = self._restore_original_order(logits, group_indices)
|
||||
if not return_dict:
|
||||
return (sorted_probs, sorted_logits)
|
||||
|
||||
return SequenceClassifierOutput(
|
||||
loss=None,
|
||||
logits=sorted_logits,
|
||||
hidden_states=out.hidden_states,
|
||||
attentions=out.attentions,
|
||||
)
|
||||
|
||||
def _reorganize_inputs(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
attention_mask: torch.Tensor,
|
||||
token_type_ids: Optional[torch.Tensor],
|
||||
) -> tuple[
|
||||
torch.Tensor, torch.Tensor, Optional[torch.Tensor], List[int], List[List[int]]
|
||||
]:
|
||||
"""
|
||||
Group inputs by unique queries: for each query, produce [query] + its passages,
|
||||
then flatten, pad, and return pair sizes and original indices mapping.
|
||||
"""
|
||||
batch_size = input_ids.size(0)
|
||||
# Structure: query_key -> {
|
||||
# 'query': (seq, mask, tt),
|
||||
# 'passages': [(seq, mask, tt), ...],
|
||||
# 'indices': [original_index, ...]
|
||||
# }
|
||||
grouped = {}
|
||||
|
||||
for idx in range(batch_size):
|
||||
seq = input_ids[idx]
|
||||
mask = attention_mask[idx]
|
||||
token_type_ids[idx] if token_type_ids is not None else torch.zeros_like(seq)
|
||||
|
||||
sep_idxs = (seq == self.config.sep_token_id).nonzero(as_tuple=True)[0]
|
||||
if sep_idxs.numel() == 0:
|
||||
raise ValueError(f"No SEP in sequence {idx}")
|
||||
first_sep = sep_idxs[0].item()
|
||||
second_sep = sep_idxs[1].item()
|
||||
|
||||
# Extract query and passage
|
||||
q_seq = seq[: first_sep + 1]
|
||||
q_mask = mask[: first_sep + 1]
|
||||
q_tt = torch.zeros_like(q_seq)
|
||||
|
||||
p_seq = seq[first_sep : second_sep + 1]
|
||||
p_mask = mask[first_sep : second_sep + 1]
|
||||
p_seq = p_seq.clone()
|
||||
p_seq[0] = self.config.cls_token_id
|
||||
p_tt = torch.zeros_like(p_seq)
|
||||
|
||||
# Build key excluding CLS/SEP
|
||||
key = tuple(
|
||||
q_seq[
|
||||
(q_seq != self.config.cls_token_id)
|
||||
& (q_seq != self.config.sep_token_id)
|
||||
].tolist()
|
||||
)
|
||||
|
||||
# truncation
|
||||
q_seq = q_seq[: self.config.max_position_embeddings]
|
||||
q_seq[-1] = self.config.sep_token_id
|
||||
p_seq = p_seq[: self.config.max_position_embeddings]
|
||||
p_seq[-1] = self.config.sep_token_id
|
||||
q_mask = q_mask[: self.config.max_position_embeddings]
|
||||
p_mask = p_mask[: self.config.max_position_embeddings]
|
||||
q_tt = q_tt[: self.config.max_position_embeddings]
|
||||
p_tt = p_tt[: self.config.max_position_embeddings]
|
||||
|
||||
if key not in grouped:
|
||||
grouped[key] = {
|
||||
"query": (q_seq, q_mask, q_tt),
|
||||
"passages": [],
|
||||
"indices": [],
|
||||
}
|
||||
grouped[key]["passages"].append((p_seq, p_mask, p_tt))
|
||||
grouped[key]["indices"].append(idx)
|
||||
|
||||
# Flatten according to group insertion order
|
||||
seqs, masks, tts, pair_nums, group_indices = [], [], [], [], []
|
||||
for key, data in grouped.items():
|
||||
q_seq, q_mask, q_tt = data["query"]
|
||||
passages = data["passages"]
|
||||
indices = data["indices"]
|
||||
# record sizes and original positions
|
||||
pair_nums.append(len(passages) + 1) # +1 for the query
|
||||
group_indices.append(indices)
|
||||
|
||||
# append query then its passages
|
||||
seqs.append(q_seq)
|
||||
masks.append(q_mask)
|
||||
tts.append(q_tt)
|
||||
for p_seq, p_mask, p_tt in passages:
|
||||
seqs.append(p_seq)
|
||||
masks.append(p_mask)
|
||||
tts.append(p_tt)
|
||||
|
||||
# Pad to uniform length
|
||||
max_len = max(s.size(0) for s in seqs)
|
||||
padded_seqs, padded_masks, padded_tts = [], [], []
|
||||
for s, m, t in zip(seqs, masks, tts):
|
||||
ps = torch.zeros(max_len, dtype=s.dtype, device=s.device)
|
||||
pm = torch.zeros(max_len, dtype=m.dtype, device=m.device)
|
||||
pt = torch.zeros(max_len, dtype=t.dtype, device=t.device)
|
||||
ps[: s.size(0)] = s
|
||||
pm[: m.size(0)] = m
|
||||
pt[: t.size(0)] = t
|
||||
padded_seqs.append(ps)
|
||||
padded_masks.append(pm)
|
||||
padded_tts.append(pt)
|
||||
|
||||
rid = torch.stack(padded_seqs)
|
||||
ram = torch.stack(padded_masks)
|
||||
rtt = torch.stack(padded_tts) if token_type_ids is not None else None
|
||||
|
||||
return rid, ram, rtt, pair_nums, group_indices
|
||||
|
||||
def _restore_original_order(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
group_indices: List[List[int]],
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Map flattened logits back so each original index gets its passage score.
|
||||
"""
|
||||
out = torch.zeros(logits.size(0), dtype=logits.dtype, device=logits.device)
|
||||
i = 0
|
||||
for indices in group_indices:
|
||||
for idx in indices:
|
||||
out[idx] = logits[i]
|
||||
i += 1
|
||||
return out.reshape(-1, 1)
|
||||
|
||||
def average_pooling(self, hidden_state, attention_mask):
|
||||
extended_attention_mask = (
|
||||
attention_mask.unsqueeze(-1)
|
||||
.expand(hidden_state.size())
|
||||
.to(dtype=hidden_state.dtype)
|
||||
)
|
||||
masked_hidden_state = hidden_state * extended_attention_mask
|
||||
sum_embeddings = torch.sum(masked_hidden_state, dim=1)
|
||||
sum_mask = extended_attention_mask.sum(dim=1)
|
||||
return sum_embeddings / sum_mask
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls, model_name_or_path, config: Optional[ListConRankerConfig] = None, **kwargs
|
||||
):
|
||||
model = super().from_pretrained(model_name_or_path, config=config, **kwargs)
|
||||
model.hf_model = BertModel.from_pretrained(
|
||||
model_name_or_path, config=model.config.bert_config, **kwargs
|
||||
)
|
||||
linear_path = hf_hub_download(
|
||||
repo_id = model_name_or_path,
|
||||
filename = "linear_in_embedding.pt",
|
||||
revision = "main",
|
||||
cache_dir = kwargs['cache_dir'] if 'cache_dir' in kwargs else None
|
||||
)
|
||||
list_transformer_path = hf_hub_download(
|
||||
repo_id = "ByteDance/ListConRanker",
|
||||
filename = "list_transformer.pt",
|
||||
revision = "main",
|
||||
cache_dir = kwargs['cache_dir'] if 'cache_dir' in kwargs else None
|
||||
)
|
||||
|
||||
try:
|
||||
model.linear_in_embedding.load_state_dict(torch.load(linear_path))
|
||||
model.list_transformer.load_state_dict(torch.load(list_transformer_path))
|
||||
except FileNotFoundError as e:
|
||||
raise e
|
||||
|
||||
return model
|
||||
|
||||
def multi_passage(
|
||||
self,
|
||||
sentences: List[List[str]],
|
||||
batch_size: int = 32,
|
||||
tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(
|
||||
"ByteDance/ListConRanker"
|
||||
),
|
||||
):
|
||||
"""
|
||||
Process multiple passages for each query.
|
||||
:param sentences: List of lists, where each inner list contains sentences for a query.
|
||||
:return: Tensor of logits for each passage.
|
||||
"""
|
||||
pairs = []
|
||||
for batch in sentences:
|
||||
if len(batch) < 2:
|
||||
raise ValueError("Each query must have at least one passage.")
|
||||
query = batch[0]
|
||||
passages = batch[1:]
|
||||
for passage in passages:
|
||||
pairs.append((query, passage))
|
||||
|
||||
total_batches = (len(pairs) + batch_size - 1) // batch_size
|
||||
total_logits = torch.zeros(len(pairs), dtype=torch.float, device=self.device)
|
||||
for batch in range(total_batches):
|
||||
batch_pairs = pairs[batch * batch_size : (batch + 1) * batch_size]
|
||||
inputs = tokenizer(
|
||||
batch_pairs,
|
||||
padding=True,
|
||||
truncation=False,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
for k, v in inputs.items():
|
||||
inputs[k] = v.to(self.device)
|
||||
|
||||
logits = self(**inputs)[0]
|
||||
total_logits[batch * batch_size : (batch + 1) * batch_size] = (
|
||||
logits.squeeze(1)
|
||||
)
|
||||
return total_logits.tolist()
|
||||
|
||||
def multi_passage_in_iterative_inference(
|
||||
self,
|
||||
sentences: List[str],
|
||||
stop_num: int = 20,
|
||||
decrement_rate: float = 0.2,
|
||||
min_filter_num: int = 10,
|
||||
tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(
|
||||
"ByteDance/ListConRanker"
|
||||
),
|
||||
):
|
||||
"""
|
||||
Process multiple passages for one query in iterative inference.
|
||||
:param sentences: List contains sentences for a query.
|
||||
:return: Tensor of logits for each passage.
|
||||
"""
|
||||
if stop_num < 1:
|
||||
raise ValueError("stop_num must be greater than 0")
|
||||
if decrement_rate <= 0 or decrement_rate >= 1:
|
||||
raise ValueError("decrement_rate must be in (0, 1)")
|
||||
if min_filter_num < 1:
|
||||
raise ValueError("min_filter_num must be greater than 0")
|
||||
|
||||
query = sentences[0]
|
||||
passage = sentences[1:]
|
||||
|
||||
filter_times = 0
|
||||
passage2score = defaultdict(list)
|
||||
while len(passage) > stop_num:
|
||||
batch = [[query] + passage]
|
||||
pred_scores = self.multi_passage(
|
||||
batch, batch_size=len(batch[0]) - 1, tokenizer=tokenizer
|
||||
)
|
||||
pred_scores_argsort = np.argsort(
|
||||
pred_scores
|
||||
).tolist() # Sort in increasing order
|
||||
|
||||
passage_len = len(passage)
|
||||
to_filter_num = math.ceil(passage_len * decrement_rate)
|
||||
if to_filter_num < min_filter_num:
|
||||
to_filter_num = min_filter_num
|
||||
|
||||
have_filter_num = 0
|
||||
while have_filter_num < to_filter_num:
|
||||
idx = pred_scores_argsort[have_filter_num]
|
||||
passage2score[passage[idx]].append(pred_scores[idx] + filter_times)
|
||||
have_filter_num += 1
|
||||
while (
|
||||
pred_scores[pred_scores_argsort[have_filter_num - 1]]
|
||||
== pred_scores[pred_scores_argsort[have_filter_num]]
|
||||
):
|
||||
idx = pred_scores_argsort[have_filter_num]
|
||||
passage2score[passage[idx]].append(pred_scores[idx] + filter_times)
|
||||
have_filter_num += 1
|
||||
next_passage = []
|
||||
next_passage_idx = have_filter_num
|
||||
while next_passage_idx < len(passage):
|
||||
idx = pred_scores_argsort[next_passage_idx]
|
||||
next_passage.append(passage[idx])
|
||||
next_passage_idx += 1
|
||||
passage = next_passage
|
||||
filter_times += 1
|
||||
|
||||
batch = [[query] + passage]
|
||||
pred_scores = self.multi_passage(
|
||||
batch, batch_size=len(batch[0]) - 1, tokenizer=tokenizer
|
||||
)
|
||||
|
||||
cnt = 0
|
||||
while cnt < len(passage):
|
||||
passage2score[passage[cnt]].append(pred_scores[cnt] + filter_times)
|
||||
cnt += 1
|
||||
|
||||
passage = sentences[1:]
|
||||
final_score = []
|
||||
for i in range(len(passage)):
|
||||
p = passage[i]
|
||||
final_score.append(passage2score[p][0])
|
||||
return final_score
|
||||
Reference in New Issue
Block a user