import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForTokenClassification
from huggingface_hub import hf_hub_download
from modeling import IENLIModel, spans_overlap
from jinja2 import Template
import spaces
model = None
tokenizer = None
device = None
def load_model():
global model, tokenizer, device
if model is not None:
return
device = torch.device("cuda")
encoder = AutoModelForTokenClassification.from_pretrained("microsoft/deberta-v3-large")
tokenizer = AutoTokenizer.from_pretrained("microsoft/deberta-v3-large")
model = IENLIModel(encoder, tokenizer)
ckpt_path = hf_hub_download(
repo_id="nicpopovic/jedi_d6a5f1cb",
repo_type="dataset",
filename="weights_d6a5f1cb.ckpt",
)
state_dict = torch.load(ckpt_path, map_location="cpu")
model.load_state_dict(state_dict)
model.to(device)
model.eval()
load_model() # <-- add this, run once at startup, before demo.launch()
# NLI label mapping
label_map = {
0: "entailment",
1: "neutral",
2: "contradiction"
}
# Process example
def process_example(premise, hypothesis):
text = premise + "[SEP]" + hypothesis
tokenized = tokenizer(text, return_offsets_mapping=True, truncation=True)
sep_token_id = tokenizer("[SEP]", add_special_tokens=False)['input_ids'][0]
sep_location = tokenized["input_ids"].index(sep_token_id)
return {
"input_ids": torch.nn.utils.rnn.pad_sequence([torch.tensor(tokenized["input_ids"])], batch_first=True, padding_value=0).to(device),
"attention_mask": torch.nn.utils.rnn.pad_sequence([torch.tensor(tokenized["attention_mask"])], batch_first=True, padding_value=0).to(device),
"offset_mapping": torch.nn.utils.rnn.pad_sequence([torch.tensor(tokenized["offset_mapping"], dtype=torch.long)], batch_first=True, padding_value=0).to(device),
"sep_locations": [sep_location],
"spans": [[]],
"spans_salient": [[]],
"positives_left": [[]],
"negatives_left": [[]],
"negatives_spans": [[]],
"labels_salient": [[]]
}
# ==== Visualization utilities ====
COLOR_PALETTE = [
'#FFB3BA', '#FFDFBA', '#FFFFBA', '#BAFFC9', '#BAE1FF',
'#E2BAFF', '#FFD1DC', '#D5AAFF', '#A0E7E5', '#B5ead7',
]
HTML_TEMPLATE = """
Data Visualization
{{ verdict }}
{{ highlighted_premise|safe }}
"""
def resolve_overlaps(spans):
kept = []
for s_start, s_end, label, prob, is_salient in spans:
overlap = False
for k_start, k_end, *_ in kept:
if not (s_end <= k_start or s_start >= k_end):
overlap = True
break
if not overlap:
kept.append((s_start, s_end, label, prob, is_salient))
return kept
def highlight_spans_flat(text, spans):
spans = sorted(spans, key=lambda x: x[0])
result = []
prev_idx = 0
for start, end, label, prob, is_salient in spans:
if start > prev_idx:
result.append(text[prev_idx:start])
if label == 0:
color = '#BAFFC9'
elif label == 2:
color = '#FFB3BA'
else:
color = 'rgb(240,240,240)'
snippet = text[start:end]
html = f'{snippet}'
if is_salient and prob is not None:
html += f'{prob:.0%}'
html += ''
result.append(html)
prev_idx = end
if prev_idx < len(text):
result.append(text[prev_idx:])
return ''.join(result)
@spaces.GPU
def predict_nli(premise, hypothesis):
input_data = process_example(premise, hypothesis)
with torch.no_grad():
predicted_label, (spans, classification, _), softmaxes, span_probs = model(input_data)
label_text = label_map.get(predicted_label[0], "Unknown")
spans_for_example = spans[0]
fact_labels = {i: int(classification[0][i].item()) for i in range(len(spans_for_example))}
salient_fact_indices = [i for i, lab in fact_labels.items() if lab in [0, 2]]
all_spans = []
salient_spans = []
for i in salient_fact_indices:
s, e = spans_for_example[i]
lab = fact_labels[i]
if predicted_label[0] == 1:
lab = 1
prob = span_probs[0][i].item()
salient = False
else:
prob = softmaxes[0][i].item() * span_probs[0][i].item()
salient = True
salient_spans.append((s, e, lab, prob, salient))
salient_spans.sort(key=lambda x: -x[3])
all_spans.extend(salient_spans)
non_salient_indices = [i for i in range(len(spans_for_example)) if i not in salient_fact_indices]
non_salient_spans = []
for i in non_salient_indices:
s, e = spans_for_example[i]
lab = fact_labels[i]
prob = span_probs[0][i].item()
non_salient_spans.append((s, e, lab, prob, False))
non_salient_spans.sort(key=lambda x: -x[3])
all_spans.extend(non_salient_spans)
final_spans = resolve_overlaps(all_spans)
highlighted_premise = highlight_spans_flat(premise, final_spans)
if label_text == "entailment":
label_text = "Entailment ✅"
elif label_text == "contradiction":
label_text = "Contradiction ❌"
else:
label_text = "Neutral ➖"
html_output = Template(HTML_TEMPLATE).render(
premise=premise,
hypothesis=hypothesis,
highlighted_premise=highlighted_premise,
verdict=label_text.capitalize()
)
return html_output.strip()
examples = [
["The Ottawa Sun is a daily tabloid newspaper in Ottawa, Ontario, Canada. It is published by Sun Media. It was first published in 1983 as the “Ottawa Sunday Herald”, until it was acquired by (then) Toronto Sun Publishing Corporation in 1988. In April 2015, Sun Media papers were acquired by Postmedia.",
"Toronto Sun Publishing acquired the Ottawa Sun in the late nineties"],
["The Washington Nationals are a professional baseball team that has been based in Washington, D.C. since . The Nationals are a member of both the Major League Baseball's (MLB) National League Eastern Division and the National League (NL) itself. Since the 2008 season, the Nationals have played in Nationals Park; from 2005 through , the team played in Robert F. Kennedy Memorial Stadium.",
"The Washington Nationals have played in Nationals Park for more than 1000 days."]
]
example_0_html = """
Contradiction ❌
The Ottawa Sun is a daily tabloid newspaper in Ottawa, Ontario, Canada. It is published by Sun Media. It was first published in 1983 as the “Ottawa Sunday Herald”, until it was acquired by (then) Toronto Sun Publishing Corporation in 1988.100% In April 2015, Sun Media papers were acquired by Postmedia.