wehub-resource-sync
593 lines
19 KiB
Python
593 lines
19 KiB
Python
# mypy: check_untyped_defs = False
|
|
###############################################
|
|
# Source: https://github.com/tingofurro/summac
|
|
###############################################
|
|
|
|
import nltk
|
|
import os
|
|
import json
|
|
import torch
|
|
from deepeval import utils as utils_misc
|
|
|
|
model_map = {
|
|
"snli-base": {
|
|
"model_card": "boychaboy/SNLI_roberta-base",
|
|
"entailment_idx": 0,
|
|
"contradiction_idx": 2,
|
|
},
|
|
"snli-large": {
|
|
"model_card": "boychaboy/SNLI_roberta-large",
|
|
"entailment_idx": 0,
|
|
"contradiction_idx": 2,
|
|
},
|
|
"mnli-base": {
|
|
"model_card": "microsoft/deberta-base-mnli",
|
|
"entailment_idx": 2,
|
|
"contradiction_idx": 0,
|
|
},
|
|
"mnli": {
|
|
"model_card": "roberta-large-mnli",
|
|
"entailment_idx": 2,
|
|
"contradiction_idx": 0,
|
|
},
|
|
"anli": {
|
|
"model_card": "ynie/roberta-large-snli_mnli_fever_anli_R1_R2_R3-nli",
|
|
"entailment_idx": 0,
|
|
"contradiction_idx": 2,
|
|
},
|
|
"vitc-base": {
|
|
"model_card": "tals/albert-base-vitaminc-mnli",
|
|
"entailment_idx": 0,
|
|
"contradiction_idx": 1,
|
|
},
|
|
"vitc": {
|
|
"model_card": "tals/albert-xlarge-vitaminc-mnli",
|
|
"entailment_idx": 0,
|
|
"contradiction_idx": 1,
|
|
},
|
|
"vitc-only": {
|
|
"model_card": "tals/albert-xlarge-vitaminc",
|
|
"entailment_idx": 0,
|
|
"contradiction_idx": 1,
|
|
},
|
|
}
|
|
|
|
|
|
def card_to_name(card):
|
|
card2name = {v["model_card"]: k for k, v in model_map.items()}
|
|
if card in card2name:
|
|
return card2name[card]
|
|
return card
|
|
|
|
|
|
def name_to_card(name):
|
|
if name in model_map:
|
|
return model_map[name]["model_card"]
|
|
return name
|
|
|
|
|
|
def get_neutral_idx(ent_idx, con_idx):
|
|
return list(set([0, 1, 2]) - set([ent_idx, con_idx]))[0]
|
|
|
|
|
|
class _SummaCImager:
|
|
def __init__(
|
|
self,
|
|
model_name="mnli",
|
|
granularity="paragraph",
|
|
use_cache=True,
|
|
max_doc_sents=100,
|
|
device="cuda",
|
|
**kwargs
|
|
):
|
|
self.grans = granularity.split("-")
|
|
|
|
assert (
|
|
all(
|
|
gran in ["paragraph", "sentence", "document", "2sents", "mixed"]
|
|
for gran in self.grans
|
|
)
|
|
and len(self.grans) <= 2
|
|
), "Unrecognized `granularity` %s" % (granularity)
|
|
assert (
|
|
model_name in model_map.keys()
|
|
), "Unrecognized model name: `%s`" % (model_name)
|
|
|
|
self.model_name = model_name
|
|
if model_name != "decomp":
|
|
self.model_card = name_to_card(model_name)
|
|
self.entailment_idx = model_map[model_name]["entailment_idx"]
|
|
self.contradiction_idx = model_map[model_name]["contradiction_idx"]
|
|
self.neutral_idx = get_neutral_idx(
|
|
self.entailment_idx, self.contradiction_idx
|
|
)
|
|
|
|
self.granularity = granularity
|
|
self.use_cache = use_cache
|
|
self.cache_folder = "/export/share/plaban/summac_cache/"
|
|
|
|
self.max_doc_sents = max_doc_sents
|
|
self.max_input_length = 500
|
|
self.device = device
|
|
self.cache = {}
|
|
self.model = None # Lazy loader
|
|
|
|
def load_nli(self):
|
|
if self.model_name == "decomp":
|
|
try:
|
|
from allennlp.predictors.predictor import Predictor
|
|
except ModuleNotFoundError:
|
|
print(
|
|
"allennlp library is not installed. "
|
|
"Please install the library by following the instruction from their documentation:"
|
|
"https://docs.allennlp.org/main/"
|
|
)
|
|
self.model = Predictor.from_path(
|
|
"https://storage.googleapis.com/allennlp-public-models/decomposable-attention-elmo-2020.04.09.tar.gz",
|
|
cuda_device=0,
|
|
)
|
|
|
|
else:
|
|
try:
|
|
from transformers import (
|
|
AutoTokenizer,
|
|
AutoModelForSequenceClassification,
|
|
)
|
|
except ModuleNotFoundError:
|
|
print(
|
|
"transformers library is not installed. Run 'pip install transformers'"
|
|
)
|
|
self.tokenizer = AutoTokenizer.from_pretrained(self.model_card)
|
|
self.model = AutoModelForSequenceClassification.from_pretrained(
|
|
self.model_card
|
|
).eval()
|
|
self.model.to(self.device)
|
|
|
|
def split_sentences(self, text):
|
|
sentences = nltk.tokenize.sent_tokenize(text)
|
|
sentences = [sent for sent in sentences if len(sent) > 10]
|
|
return sentences
|
|
|
|
def split_2sents(self, text):
|
|
sentences = nltk.tokenize.sent_tokenize(text)
|
|
sentences = [sent for sent in sentences if len(sent) > 10]
|
|
two_sents = [
|
|
" ".join(sentences[i : (i + 2)]) for i in range(len(sentences))
|
|
]
|
|
return two_sents
|
|
|
|
def split_paragraphs(self, text):
|
|
if text.count("\n\n") > 0:
|
|
paragraphs = [p.strip() for p in text.split("\n\n")]
|
|
else:
|
|
paragraphs = [p.strip() for p in text.split("\n")]
|
|
return [p for p in paragraphs if len(p) > 10]
|
|
|
|
def split_text(self, text, granularity="sentence"):
|
|
if granularity == "document":
|
|
return [text]
|
|
elif granularity == "paragraph":
|
|
return self.split_paragraphs(text)
|
|
elif granularity == "sentence":
|
|
return self.split_sentences(text)
|
|
elif granularity == "2sents":
|
|
return self.split_2sents(text)
|
|
elif granularity == "mixed":
|
|
return self.split_sentences(text) + self.split_paragraphs(text)
|
|
|
|
def build_image(self, original, generated):
|
|
import numpy as np
|
|
|
|
cache_key = (original, generated)
|
|
if self.use_cache and cache_key in self.cache:
|
|
cached_image = self.cache[cache_key]
|
|
cached_image = cached_image[:, : self.max_doc_sents, :]
|
|
return cached_image
|
|
|
|
if len(self.grans) == 1:
|
|
gran_doc, gran_sum = self.grans[0], self.grans[0]
|
|
else:
|
|
gran_doc, gran_sum = self.grans[0], self.grans[1]
|
|
|
|
original_chunks = self.split_text(original, granularity=gran_doc)[
|
|
: self.max_doc_sents
|
|
]
|
|
generated_chunks = self.split_text(generated, granularity=gran_sum)
|
|
|
|
N_ori = len(original_chunks)
|
|
N_gen = len(generated_chunks)
|
|
|
|
if N_ori == 0 or N_gen == 0:
|
|
return np.zeros((3, 1, 1))
|
|
# assert (N_ori > 0 and N_gen > 0), "One of the inputs has no chunks"
|
|
|
|
image = np.zeros((3, N_ori, N_gen))
|
|
|
|
if self.model is None:
|
|
self.load_nli()
|
|
|
|
dataset = [
|
|
{
|
|
"premise": original_chunks[i],
|
|
"hypothesis": generated_chunks[j],
|
|
"doc_i": i,
|
|
"gen_i": j,
|
|
}
|
|
for i in range(N_ori)
|
|
for j in range(N_gen)
|
|
]
|
|
for batch in utils_misc.batcher(dataset, batch_size=20):
|
|
if self.model_name == "decomp":
|
|
batch_evids, batch_conts, batch_neuts = [], [], []
|
|
batch_json = [
|
|
{"premise": d["premise"], "hypothesis": d["hypothesis"]}
|
|
for d in batch
|
|
]
|
|
model_outs = self.model.predict_batch_json(batch_json)
|
|
for out in model_outs:
|
|
probs = out["label_probs"]
|
|
batch_evids.append(probs[0])
|
|
batch_conts.append(probs[1])
|
|
batch_neuts.append(probs[2])
|
|
|
|
else:
|
|
batch_prems = [b["premise"] for b in batch]
|
|
batch_hypos = [b["hypothesis"] for b in batch]
|
|
batch_tokens = self.tokenizer.batch_encode_plus(
|
|
list(zip(batch_prems, batch_hypos)),
|
|
padding=True,
|
|
truncation=True,
|
|
max_length=self.max_input_length,
|
|
return_tensors="pt",
|
|
truncation_strategy="only_first",
|
|
)
|
|
batch_tokens = {
|
|
k: v.to(self.device) for k, v in batch_tokens.items()
|
|
}
|
|
with torch.no_grad():
|
|
model_outputs = self.model(**batch_tokens)
|
|
|
|
batch_probs = torch.nn.functional.softmax(
|
|
model_outputs["logits"], dim=-1
|
|
)
|
|
batch_evids = batch_probs[:, self.entailment_idx].tolist()
|
|
batch_conts = batch_probs[:, self.contradiction_idx].tolist()
|
|
batch_neuts = batch_probs[:, self.neutral_idx].tolist()
|
|
|
|
for b, evid, cont, neut in zip(
|
|
batch, batch_evids, batch_conts, batch_neuts
|
|
):
|
|
image[0, b["doc_i"], b["gen_i"]] = evid
|
|
image[1, b["doc_i"], b["gen_i"]] = cont
|
|
image[2, b["doc_i"], b["gen_i"]] = neut
|
|
|
|
if self.use_cache:
|
|
self.cache[cache_key] = image
|
|
return image
|
|
|
|
def get_cache_file(self):
|
|
return os.path.join(
|
|
self.cache_folder,
|
|
"cache_%s_%s.json" % (self.model_name, self.granularity),
|
|
)
|
|
|
|
def save_cache(self):
|
|
cache_cp = {"[///]".join(k): v.tolist() for k, v in self.cache.items()}
|
|
with open(self.get_cache_file(), "w") as f:
|
|
json.dump(cache_cp, f)
|
|
|
|
def load_cache(self):
|
|
import numpy as np
|
|
|
|
cache_file = self.get_cache_file()
|
|
if os.path.isfile(cache_file):
|
|
with open(cache_file, "r") as f:
|
|
cache_cp = json.load(f)
|
|
self.cache = {
|
|
tuple(k.split("[///]")): np.array(v)
|
|
for k, v in cache_cp.items()
|
|
}
|
|
|
|
|
|
class _SummaCConv(torch.nn.Module):
|
|
def __init__(
|
|
self,
|
|
models=["mnli", "anli", "vitc"],
|
|
bins="even50",
|
|
granularity="sentence",
|
|
nli_labels="e",
|
|
device="cuda",
|
|
start_file=None,
|
|
imager_load_cache=True,
|
|
agg="mean",
|
|
norm_histo=False,
|
|
**kwargs
|
|
):
|
|
import numpy as np
|
|
|
|
# `bins` should be `even%d` or `percentiles`
|
|
assert nli_labels in [
|
|
"e",
|
|
"c",
|
|
"n",
|
|
"ec",
|
|
"en",
|
|
"cn",
|
|
"ecn",
|
|
], "Unrecognized nli_labels argument %s" % (nli_labels)
|
|
|
|
super(SummaCConv, self).__init__()
|
|
self.device = device
|
|
self.models = models
|
|
|
|
self.imagers = []
|
|
for model_name in models:
|
|
self.imagers.append(
|
|
SummaCImager(
|
|
model_name=model_name, granularity=granularity, **kwargs
|
|
)
|
|
)
|
|
if imager_load_cache:
|
|
for imager in self.imagers:
|
|
imager.load_cache()
|
|
assert len(self.imagers) > 0, "Imager names were empty or unrecognized"
|
|
|
|
if "even" in bins:
|
|
n_bins = int(bins.replace("even", ""))
|
|
self.bins = list(np.arange(0, 1, 1 / n_bins)) + [1.0]
|
|
elif bins == "percentile":
|
|
self.bins = [
|
|
0.0,
|
|
0.01,
|
|
0.02,
|
|
0.03,
|
|
0.04,
|
|
0.07,
|
|
0.13,
|
|
0.37,
|
|
0.90,
|
|
0.91,
|
|
0.92,
|
|
0.93,
|
|
0.94,
|
|
0.95,
|
|
0.955,
|
|
0.96,
|
|
0.965,
|
|
0.97,
|
|
0.975,
|
|
0.98,
|
|
0.985,
|
|
0.99,
|
|
0.995,
|
|
1.0,
|
|
]
|
|
|
|
self.nli_labels = nli_labels
|
|
self.n_bins = len(self.bins) - 1
|
|
self.norm_histo = norm_histo
|
|
self.n_rows = 10
|
|
self.n_labels = 2
|
|
self.n_depth = len(self.imagers) * len(self.nli_labels)
|
|
self.full_size = self.n_depth * self.n_bins
|
|
if self.norm_histo:
|
|
self.full_size += (
|
|
2 # Will explicitly give the count of originals and generateds
|
|
)
|
|
|
|
self.agg = agg
|
|
|
|
self.mlp = torch.nn.Linear(self.full_size, 1).to(device)
|
|
self.layer_final = torch.nn.Linear(3, self.n_labels).to(device)
|
|
|
|
if start_file is not None:
|
|
print(self.load_state_dict(torch.load(start_file)))
|
|
|
|
def build_image(self, original, generated):
|
|
import numpy as np
|
|
|
|
images = [
|
|
imager.build_image(original, generated) for imager in self.imagers
|
|
]
|
|
image = np.concatenate(images, axis=0)
|
|
return image
|
|
|
|
def compute_histogram(self, original=None, generated=None, image=None):
|
|
import numpy as np
|
|
|
|
# Takes the two texts, and generates a (n_rows, 2*n_bins)
|
|
|
|
if image is None:
|
|
image = self.build_image(original, generated)
|
|
|
|
N_depth, N_ori, N_gen = image.shape
|
|
|
|
full_histogram = []
|
|
for i_gen in range(N_gen):
|
|
histos = []
|
|
|
|
for i_depth in range(N_depth):
|
|
if (
|
|
(i_depth % 3 == 0 and "e" in self.nli_labels)
|
|
or (i_depth % 3 == 1 and "c" in self.nli_labels)
|
|
or (i_depth % 3 == 2 and "n" in self.nli_labels)
|
|
):
|
|
histo, X = np.histogram(
|
|
image[i_depth, :, i_gen],
|
|
range=(0, 1),
|
|
bins=self.bins,
|
|
density=self.norm_histo,
|
|
)
|
|
histos.append(histo)
|
|
|
|
if self.norm_histo:
|
|
histos = [[N_ori, N_gen]] + histos
|
|
histogram_row = np.concatenate(histos)
|
|
full_histogram.append(histogram_row)
|
|
|
|
n_rows_missing = self.n_rows - len(full_histogram)
|
|
full_histogram += [[0.0] * self.full_size] * n_rows_missing
|
|
full_histogram = full_histogram[: self.n_rows]
|
|
full_histogram = np.array(full_histogram)
|
|
return image, full_histogram
|
|
|
|
def forward(self, originals, generateds, images=None):
|
|
if images is not None:
|
|
# In case they've been pre-computed.
|
|
histograms = []
|
|
for image in images:
|
|
_, histogram = self.compute_histogram(image=image)
|
|
histograms.append(histogram)
|
|
else:
|
|
images, histograms = [], []
|
|
for original, generated in zip(originals, generateds):
|
|
image, histogram = self.compute_histogram(
|
|
original=original, generated=generated
|
|
)
|
|
images.append(image)
|
|
histograms.append(histogram)
|
|
|
|
N = len(histograms)
|
|
histograms = torch.FloatTensor(histograms).to(self.device)
|
|
|
|
non_zeros = (torch.sum(histograms, dim=-1) != 0.0).long()
|
|
seq_lengths = non_zeros.sum(dim=-1).tolist()
|
|
|
|
mlp_outs = self.mlp(histograms).reshape(N, self.n_rows)
|
|
features = []
|
|
|
|
for mlp_out, seq_length in zip(mlp_outs, seq_lengths):
|
|
if seq_length > 0:
|
|
Rs = mlp_out[:seq_length]
|
|
if self.agg == "mean":
|
|
features.append(
|
|
torch.cat(
|
|
[
|
|
torch.mean(Rs).unsqueeze(0),
|
|
torch.mean(Rs).unsqueeze(0),
|
|
torch.mean(Rs).unsqueeze(0),
|
|
]
|
|
).unsqueeze(0)
|
|
)
|
|
elif self.agg == "min":
|
|
features.append(
|
|
torch.cat(
|
|
[
|
|
torch.min(Rs).unsqueeze(0),
|
|
torch.min(Rs).unsqueeze(0),
|
|
torch.min(Rs).unsqueeze(0),
|
|
]
|
|
).unsqueeze(0)
|
|
)
|
|
elif self.agg == "max":
|
|
features.append(
|
|
torch.cat(
|
|
[
|
|
torch.max(Rs).unsqueeze(0),
|
|
torch.max(Rs).unsqueeze(0),
|
|
torch.max(Rs).unsqueeze(0),
|
|
]
|
|
).unsqueeze(0)
|
|
)
|
|
elif self.agg == "all":
|
|
features.append(
|
|
torch.cat(
|
|
[
|
|
torch.min(Rs).unsqueeze(0),
|
|
torch.mean(Rs).unsqueeze(0),
|
|
torch.max(Rs).unsqueeze(0),
|
|
]
|
|
).unsqueeze(0)
|
|
)
|
|
else:
|
|
features.append(
|
|
torch.FloatTensor([0.0, 0.0, 0.0]).unsqueeze(0)
|
|
) # .cuda()
|
|
features = torch.cat(features)
|
|
logits = self.layer_final(features)
|
|
histograms_out = [histogram.cpu().numpy() for histogram in histograms]
|
|
return logits, histograms_out, images
|
|
|
|
def save_imager_cache(self):
|
|
for imager in self.imagers:
|
|
imager.save_cache()
|
|
|
|
def score(self, originals, generateds, **kwargs):
|
|
with torch.no_grad():
|
|
logits, histograms, images = self.forward(originals, generateds)
|
|
probs = torch.nn.functional.softmax(logits, dim=-1)
|
|
batch_scores = probs[:, 1].tolist()
|
|
return {
|
|
"scores": batch_scores
|
|
} # , "histograms": histograms, "images": images
|
|
|
|
|
|
class _SummaCZS:
|
|
def __init__(
|
|
self,
|
|
model_name="mnli",
|
|
granularity="paragraph",
|
|
op1="max",
|
|
op2="mean",
|
|
use_ent=True,
|
|
use_con=True,
|
|
imager_load_cache=True,
|
|
device="cuda",
|
|
**kwargs
|
|
):
|
|
assert op2 in ["min", "mean", "max"], "Unrecognized `op2`"
|
|
assert op1 in ["max", "mean", "min"], "Unrecognized `op1`"
|
|
|
|
self.imager = _SummaCImager(
|
|
model_name=model_name,
|
|
granularity=granularity,
|
|
device=device,
|
|
**kwargs
|
|
)
|
|
if imager_load_cache:
|
|
self.imager.load_cache()
|
|
self.op2 = op2
|
|
self.op1 = op1
|
|
self.use_ent = use_ent
|
|
self.use_con = use_con
|
|
|
|
def save_imager_cache(self):
|
|
self.imager.save_cache()
|
|
|
|
def score_one(self, original, generated):
|
|
import numpy as np
|
|
|
|
image = self.imager.build_image(original, generated)
|
|
|
|
ent_scores = np.max(image[0], axis=0)
|
|
co_scores = np.max(image[1], axis=0)
|
|
if self.op1 == "mean":
|
|
ent_scores = np.mean(image[0], axis=0)
|
|
co_scores = np.mean(image[1], axis=0)
|
|
elif self.op1 == "min":
|
|
ent_scores = np.min(image[0], axis=0)
|
|
co_scores = np.min(image[1], axis=0)
|
|
|
|
if self.use_ent and self.use_con:
|
|
scores = ent_scores - co_scores
|
|
elif self.use_ent:
|
|
scores = ent_scores
|
|
elif self.use_con:
|
|
scores = 1.0 - co_scores
|
|
|
|
final_score = np.mean(scores)
|
|
if self.op2 == "min":
|
|
final_score = np.min(scores)
|
|
elif self.op2 == "max":
|
|
final_score = np.max(scores)
|
|
|
|
return {"score": final_score, "image": image}
|
|
|
|
def score(self, sources, generateds, **kwargs):
|
|
output = {"scores": [], "images": []}
|
|
for source, gen in zip(sources, generateds):
|
|
score = self.score_one(source, gen)
|
|
output["scores"].append(score["score"])
|
|
output["images"].append(score["image"])
|
|
return output
|