Spaces:
Running on Zero
Running on Zero
| 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 = """ | |
| <!doctype html> | |
| <html lang="en"> | |
| <head> | |
| <meta charset="utf-8"> | |
| <title>Data Visualization</title> | |
| <link rel="stylesheet" type="text/css" | |
| href="https://cdn.rawgit.com/dreampulse/computer-modern-web-font/master/fonts.css"> | |
| <style> | |
| html, body { font-family: "Computer Modern Sans" sans-serif; } | |
| .highlight { border: 2px solid rgba(0, 0, 0, 0.15); background: white; padding: 2px 5px; border-radius: 5px; display: inline; } | |
| .small-text { padding: 2px 5px; border-radius: 5px; font-size: x-small; margin-left: 0.5em; vertical-align: 0.2em; font-weight: bold; font-family: monospace; } | |
| ol.no-counter { list-style: none; padding-left: 1em; } | |
| .grid-container { display: table; width: 100%;} | |
| .header { font-family: "Computer Modern Sans", sans-serif;} | |
| p { font-family: "Computer Modern Sans", sans-serif;} | |
| </style> | |
| </head> | |
| <body> | |
| <div class="grid-row"> | |
| <p><span style="font-weight: bold;">{{ verdict }}</span></p> | |
| <p style="line-height: 180%;text-align: justify;">{{ highlighted_premise|safe }}</p> | |
| </div> | |
| </div> | |
| </body> | |
| </html> | |
| """ | |
| 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'<span class="highlight" style="background-color: {color}">{snippet}' | |
| if is_salient and prob is not None: | |
| html += f'<span class="small-text" style="background-color: rgba(0,0,0,0.45); color: {color}">{prob:.0%}</span>' | |
| html += '</span>' | |
| result.append(html) | |
| prev_idx = end | |
| if prev_idx < len(text): | |
| result.append(text[prev_idx:]) | |
| return ''.join(result) | |
| 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 = """ | |
| <link rel="stylesheet" type="text/css" | |
| href="https://cdn.rawgit.com/dreampulse/computer-modern-web-font/master/fonts.css"> | |
| <style> | |
| html, body { font-family: "Computer Modern Sans" sans-serif; } | |
| .highlight { border: 2px solid rgba(0, 0, 0, 0.15); background: white; padding: 2px 5px; border-radius: 5px; display: inline; } | |
| .small-text { padding: 2px 5px; border-radius: 5px; font-size: x-small; margin-left: 0.5em; vertical-align: 0.2em; font-weight: bold; font-family: monospace; } | |
| ol.no-counter { list-style: none; padding-left: 1em; } | |
| .grid-container { display: table; width: 100%;} | |
| .header { font-family: "Computer Modern Sans", sans-serif;} | |
| p { font-family: "Computer Modern Sans", sans-serif;} | |
| </style> | |
| <p><span style="font-weight: bold;">Contradiction ❌</span></p> | |
| <p style="line-height: 180%;text-align: justify;"><span class="highlight" style="background-color: rgb(240,240,240)">The Ottawa Sun is a daily tabloid newspaper in Ottawa, Ontario, Canada.</span> It<span class="highlight" style="background-color: rgb(240,240,240)"> is published by Sun Media.</span> It<span class="highlight" style="background-color: rgb(240,240,240)"> was first published in 1983 as the “Ottawa Sunday Herald”,</span> until it was acquired by (then)<span class="highlight" style="background-color: #FFB3BA"> Toronto Sun Publishing Corporation in 1988.<span class="small-text" style="background-color: rgba(0,0,0,0.45); color: #FFB3BA">100%</span></span> In<span class="highlight" style="background-color: rgb(240,240,240)"> April 2015, Sun Media papers were acquired by Postmedia.</span></p> | |
| """ | |
| with gr.Blocks(css="footer{display:none !important} .gradio-container {padding: 0!important; height:400px;}", fill_width=True) as demo: | |
| with gr.Column(): | |
| output_display = gr.HTML(label="Highlighted Premise Spans", value=example_0_html) | |
| premise_input = gr.Textbox(label="Premise", lines=4, placeholder="Enter premise here...", value=examples[0][0]) | |
| hypothesis_input = gr.Textbox(label="Hypothesis", lines=2, placeholder="Enter hypothesis here...", value=examples[0][1]) | |
| with gr.Row(): | |
| submit_btn = gr.Button("Submit", variant="primary") | |
| examples_component = gr.Examples( | |
| examples=examples, | |
| inputs=[premise_input, hypothesis_input], | |
| outputs=output_display, | |
| fn=predict_nli, | |
| cache_examples=True, | |
| preload=0, | |
| ) | |
| submit_btn.click(fn=predict_nli, inputs=[premise_input, hypothesis_input], outputs=output_display) | |
| if __name__ == "__main__": | |
| demo.launch(ssr_mode=False) | |