|
from transformers import Pipeline |
|
import torch |
|
from typing import Union |
|
|
|
|
|
class DocumentSentenceRelevancePipeline(Pipeline): |
|
|
|
def _sanitize_parameters(self, **kwargs): |
|
threshold = kwargs.get("threshold", 0.5) |
|
return {}, {}, {"threshold": threshold} |
|
|
|
def preprocess(self, inputs): |
|
question = inputs.get("question", "") |
|
context = inputs.get("context", [""]) |
|
response = inputs.get("response", "") |
|
|
|
q_enc = self.tokenizer(question, add_special_tokens=True, truncation=False, padding=False) |
|
r_enc = self.tokenizer(response, add_special_tokens=True, truncation=False, padding=False) |
|
|
|
question_ids = q_enc["input_ids"] |
|
response_ids = r_enc["input_ids"] |
|
|
|
document_sentences_ids = [] |
|
for s in context: |
|
s_enc = self.tokenizer(s, add_special_tokens=True, truncation=False, padding=False) |
|
document_sentences_ids.append(s_enc["input_ids"]) |
|
|
|
ids = question_ids + response_ids |
|
pair_ids = [] |
|
for s_ids in document_sentences_ids: |
|
pair_ids.extend(s_ids) |
|
|
|
total_length = len(ids) + len(pair_ids) |
|
if total_length > self.tokenizer.model_max_length: |
|
num_tokens_to_remove = total_length - self.tokenizer.model_max_length |
|
ids, pair_ids, _ = self.tokenizer.truncate_sequences( |
|
ids=ids, |
|
pair_ids=pair_ids, |
|
num_tokens_to_remove=num_tokens_to_remove, |
|
truncation_strategy="only_second", |
|
stride=0, |
|
) |
|
combined_ids = ids + pair_ids |
|
token_types = [0]*len(ids) + [1]*len(pair_ids) |
|
attention_mask = [1]*len(combined_ids) |
|
|
|
sentence_positions = [] |
|
current_pos = len(ids) |
|
found_sentences = 0 |
|
|
|
for i, tok_id in enumerate(pair_ids): |
|
if tok_id == self.tokenizer.cls_token_id: |
|
sentence_positions.append(current_pos + i) |
|
found_sentences += 1 |
|
|
|
input_ids = torch.tensor([combined_ids], dtype=torch.long) |
|
attention_mask = torch.tensor([attention_mask], dtype=torch.long) |
|
token_type_ids = torch.tensor([token_types], dtype=torch.long) |
|
sentence_positions = torch.tensor([sentence_positions], dtype=torch.long) |
|
|
|
return { |
|
"input_ids": input_ids, |
|
"attention_mask": attention_mask, |
|
"token_type_ids": token_type_ids, |
|
"sentence_positions": sentence_positions |
|
} |
|
|
|
def _forward(self, model_inputs): |
|
return self.model(**model_inputs) |
|
|
|
def __call__(self, inputs: Union[dict[str, str], list[dict[str, str]]], **kwargs): |
|
if isinstance(inputs, dict): |
|
inputs = [inputs] |
|
model_outputs = super().__call__(inputs, **kwargs) |
|
pipeline_outputs = [] |
|
for i, output in enumerate(model_outputs): |
|
sentences = inputs[i]["context"] |
|
sentences_dict = { |
|
"sentence": sentences, |
|
"label": output["sentences"]["label"], |
|
"score": output["sentences"]["score"] |
|
} |
|
|
|
final_output = { |
|
"document": output["document"], |
|
"sentences": [ |
|
{ |
|
"sentence": sent, |
|
"label": label, |
|
"score": score |
|
} |
|
for sent, label, score in zip( |
|
sentences_dict["sentence"], |
|
sentences_dict["label"], |
|
sentences_dict["score"] |
|
) |
|
] |
|
} |
|
pipeline_outputs.append(final_output) |
|
return pipeline_outputs |
|
|
|
def postprocess(self, model_outputs, threshold = 0.5): |
|
doc_logits = model_outputs.doc_logits |
|
sent_logits = model_outputs.sent_logits |
|
document_probabilities = torch.softmax(doc_logits, dim=-1) |
|
sentence_probabilities = torch.softmax(sent_logits, dim=-1) |
|
|
|
document_best_class = (document_probabilities[:, 1] > threshold).long() |
|
sentence_best_class = (sentence_probabilities[:, :, 1] > threshold).long() |
|
document_score = document_probabilities[:, document_best_class] |
|
|
|
sentence_best_class = sentence_best_class.squeeze() |
|
sentence_probabilities = sentence_probabilities.squeeze() |
|
|
|
if len(sentence_best_class.shape) == 0: |
|
sentence_best_class = sentence_best_class.unsqueeze(0) |
|
sentence_probabilities = sentence_probabilities.unsqueeze(0) |
|
|
|
batch_indices = torch.arange(len(sentence_best_class)) |
|
sentence_scores = sentence_probabilities[batch_indices, sentence_best_class] |
|
|
|
best_document_label = document_best_class.numpy().item() |
|
best_document_label = self.model.config.id2label[best_document_label] |
|
|
|
best_sentence_labels = sentence_best_class.numpy().tolist() |
|
best_sentence_labels = [self.model.config.id2label[label] for label in best_sentence_labels] |
|
|
|
document_output = {"label": best_document_label, "score": document_score.numpy().item()} |
|
sentence_output = {"label": best_sentence_labels, "score": sentence_scores.numpy().tolist()} |
|
return {"document": document_output, "sentences": sentence_output} |
