|
from typing import List, Optional |
|
|
|
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig |
|
import torch |
|
from haystack.nodes.base import BaseComponent |
|
from haystack.modeling.utils import initialize_device_settings |
|
from haystack.schema import Document |
|
|
|
|
|
class EntailmentChecker(BaseComponent): |
|
""" |
|
This node checks the entailment between every document content and the query. |
|
It enrichs the documents metadata with entailment informations. |
|
It also returns aggregate entailment information. |
|
""" |
|
|
|
outgoing_edges = 1 |
|
|
|
def __init__( |
|
self, |
|
model_name_or_path: str = "roberta-large-mnli", |
|
model_version: Optional[str] = None, |
|
tokenizer: Optional[str] = None, |
|
use_gpu: bool = True, |
|
batch_size: int = 16, |
|
entailment_contradiction_threshold: float = 0.5, |
|
): |
|
""" |
|
Load a Natural Language Inference model from Transformers. |
|
|
|
:param model_name_or_path: Directory of a saved model or the name of a public model. |
|
See https://huggingface.co/models for full list of available models. |
|
:param model_version: The version of model to use from the HuggingFace model hub. Can be tag name, branch name, or commit hash. |
|
:param tokenizer: Name of the tokenizer (usually the same as model) |
|
:param use_gpu: Whether to use GPU (if available). |
|
:param batch_size: Number of Documents to be processed at a time. |
|
:param entailment_contradiction_threshold: if in the first N documents there is a strong evidence of entailment/contradiction |
|
(aggregate entailment or contradiction are greater than the threshold), the less relevant documents are not taken into account |
|
""" |
|
super().__init__() |
|
|
|
self.devices, _ = initialize_device_settings(use_cuda=use_gpu, multi_gpu=False) |
|
|
|
tokenizer = tokenizer or model_name_or_path |
|
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer) |
|
self.model = AutoModelForSequenceClassification.from_pretrained( |
|
pretrained_model_name_or_path=model_name_or_path, revision=model_version |
|
) |
|
self.batch_size = batch_size |
|
self.entailment_contradiction_threshold = entailment_contradiction_threshold |
|
self.model.to(str(self.devices[0])) |
|
|
|
id2label = AutoConfig.from_pretrained(model_name_or_path).id2label |
|
self.labels = [id2label[k].lower() for k in sorted(id2label)] |
|
if "entailment" not in self.labels: |
|
raise ValueError( |
|
"The model config must contain entailment value in the id2label dict." |
|
) |
|
|
|
def run(self, query: str, documents: List[Document]): |
|
|
|
scores, agg_con, agg_neu, agg_ent = 0, 0, 0, 0 |
|
for i, doc in enumerate(documents): |
|
entailment_info = self.get_entailment(premise=doc.content, hypotesis=query) |
|
doc.meta["entailment_info"] = entailment_info |
|
|
|
scores += doc.score |
|
con, neu, ent = ( |
|
entailment_info["contradiction"], |
|
entailment_info["neutral"], |
|
entailment_info["entailment"], |
|
) |
|
agg_con += con * doc.score |
|
agg_neu += neu * doc.score |
|
agg_ent += ent * doc.score |
|
|
|
|
|
|
|
if max(agg_con, agg_ent) / scores > self.entailment_contradiction_threshold: |
|
break |
|
|
|
aggregate_entailment_info = { |
|
"contradiction": round(agg_con / scores, 2), |
|
"neutral": round(agg_neu / scores, 2), |
|
"entailment": round(agg_ent / scores, 2), |
|
} |
|
|
|
entailment_checker_result = { |
|
"documents": documents[: i + 1], |
|
"aggregate_entailment_info": aggregate_entailment_info, |
|
} |
|
|
|
return entailment_checker_result, "output_1" |
|
|
|
def run_batch(self, queries: List[str], documents: List[Document]): |
|
pass |
|
|
|
def get_entailment(self, premise, hypotesis): |
|
with torch.inference_mode(): |
|
inputs = self.tokenizer( |
|
f"{premise}{self.tokenizer.sep_token}{hypotesis}", return_tensors="pt" |
|
).to(self.devices[0]) |
|
out = self.model(**inputs) |
|
logits = out.logits |
|
probs = ( |
|
torch.nn.functional.softmax(logits, dim=-1)[0, :].detach().cpu().numpy() |
|
) |
|
entailment_dict = {k.lower(): v for k, v in zip(self.labels, probs)} |
|
return entailment_dict |
|
|
