kmfoda's picture
Fix columns error
744f7bb
import argparse
import json
import operator
import os
import re
from pathlib import Path
import spacy
import spacy.lang.en
import streamlit as st
from meerkat import DataPanel
from spacy.tokens import Doc
from align import NGramAligner, BertscoreAligner, StaticEmbeddingAligner
from components import MainView
from utils import clean_text
MIN_SEMANTIC_SIM_THRESHOLD = 0.1
MAX_SEMANTIC_SIM_TOP_K = 10
Doc.set_extension("name", default=None, force=True)
Doc.set_extension("column", default=None, force=True)
class Instance():
def __init__(self, id_, document, reference, preds, data=None):
self.id = id_
self.document = document
self.reference = reference
self.preds = preds
self.data = data
@st.cache(allow_output_mutation=True)
def load_from_index(filename, index):
with open(filename) as f:
for i, line in enumerate(f):
if i == index:
return json.loads(line.strip())
def _nlp_key(x: spacy.Language):
return str(x.path)
@st.cache(allow_output_mutation=True, hash_funcs={spacy.lang.en.English: _nlp_key})
def load_dataset(path: str, nlp: spacy.Language):
if path.endswith('.jsonl'):
return DataPanel.from_jsonl(path)
try:
return DataPanel.read(path, nlp=nlp)
except NotADirectoryError:
return DataPanel.from_jsonl(path)
@st.cache(allow_output_mutation=True)
def get_nlp():
try:
nlp = spacy.load("en_core_web_lg")
except:
nlp = spacy.load("en_core_web_sm")
is_lg = False
else:
is_lg = True
nlp.add_pipe('sentencizer', before="parser")
return nlp, is_lg
def retrieve(dataset, index, filename=None):
if index >= len(dataset):
st.error(f"Index {index} exceeds dataset length.")
eval_dataset = None
if filename:
# TODO Handle this through dedicated fields
if "cnn_dailymail" in filename:
eval_dataset = "cnndm"
elif "xsum" in filename:
eval_dataset = "xsum"
data = dataset[index]
id_ = data.get('id', '')
try:
document = data['spacy:document']
except KeyError:
if not is_lg:
st.error("'en_core_web_lg model' is required unless loading from cached file."
"To install: 'python -m spacy download en_core_web_lg'")
try:
text = data['document']
except KeyError:
text = data['article']
if not text:
st.error("Document is blank")
return
document = nlp(text if args.no_clean else clean_text(text))
document._.name = "Document"
document._.column = "document"
try:
reference = data['spacy:summary:reference']
except KeyError:
if not is_lg:
st.error("'en_core_web_lg model' is required unless loading from cached file."
"To install: 'python -m spacy download en_core_web_lg'")
try:
text = data['summary'] if 'summary' in data else data['summary:reference']
except KeyError:
text = data.get('highlights')
if text:
reference = nlp(text if args.no_clean else clean_text(text))
else:
reference = None
if reference is not None:
reference._.name = "Reference"
reference._.column = "summary:reference"
model_names = set()
for k in data:
m = re.match('(preprocessed_)?summary:(?P<model>.*)', k)
if m:
model_name = m.group('model')
if model_name != 'reference':
model_names.add(model_name)
preds = []
for model_name in model_names:
try:
pred = data[f"spacy:summary:{model_name}"]
except KeyError:
if not is_lg:
st.error("'en_core_web_lg model' is required unless loading from cached file."
"To install: 'python -m spacy download en_core_web_lg'")
text = data[f"summary:{model_name}"]
pred = nlp(text if args.no_clean else clean_text(text))
parts = model_name.split("-")
primary_sort = 0
if len(parts) == 2:
model, train_dataset = parts
if train_dataset == eval_dataset:
formatted_model_name = model.upper()
else:
formatted_model_name = f"{model.upper()} ({train_dataset.upper()}-trained)"
if train_dataset in ["xsum", "cnndm"]:
primary_sort = 1
else:
primary_sort = 2
else:
formatted_model_name = model_name.upper()
pred._.name = formatted_model_name
pred._.column = f"summary:{model_name}"
preds.append(
((primary_sort, formatted_model_name), pred)
)
preds = [pred for _, pred in sorted(preds)]
return Instance(
id_=id_,
document=document,
reference=reference,
preds=preds,
data=data,
)
def filter_alignment(alignment, threshold, top_k):
filtered_alignment = {}
for k, v in alignment.items():
filtered_matches = [(match_idx, score) for match_idx, score in v if score >= threshold]
if filtered_matches:
filtered_alignment[k] = sorted(filtered_matches, key=operator.itemgetter(1), reverse=True)[:top_k]
return filtered_alignment
def select_comparison(example):
all_summaries = []
if example.reference:
all_summaries.append(example.reference)
if example.preds:
all_summaries.extend(example.preds)
from_documents = [example.document]
if example.reference:
from_documents.append(example.reference)
document_names = [document._.name for document in from_documents]
select_document_name = sidebar_placeholder_from.selectbox(
label="Comparison FROM:",
options=document_names
)
document_index = document_names.index(select_document_name)
selected_document = from_documents[document_index]
remaining_summaries = [summary for summary in all_summaries if
summary._.name != selected_document._.name]
remaining_summary_names = [summary._.name for summary in remaining_summaries]
selected_summary_names = sidebar_placeholder_to.multiselect(
'Comparison TO:',
remaining_summary_names,
remaining_summary_names
)
selected_summaries = []
for summary_name in selected_summary_names:
summary_index = remaining_summary_names.index(summary_name)
selected_summaries.append(remaining_summaries[summary_index])
return selected_document, selected_summaries
def show_main(example):
# Get user input
semantic_sim_type = st.sidebar.radio(
"Semantic similarity type:",
["Contextual embedding", "Static embedding"]
)
semantic_sim_threshold = st.sidebar.slider(
"Semantic similarity threshold:",
min_value=MIN_SEMANTIC_SIM_THRESHOLD,
max_value=1.0,
step=0.1,
value=0.2,
)
semantic_sim_top_k = st.sidebar.slider(
"Semantic similarity top-k:",
min_value=1,
max_value=MAX_SEMANTIC_SIM_TOP_K,
step=1,
value=10,
)
document, summaries = select_comparison(example)
layout = st.sidebar.radio("Layout:", ["Vertical", "Horizontal"]).lower()
scroll = True
gray_out_stopwords = st.sidebar.checkbox(label="Gray out stopwords", value=True)
# Gather data
try:
lexical_alignments = [
example.data[f'{NGramAligner.__name__}:spacy:{document._.column}:spacy:{summary._.column}']
for summary in summaries
]
except KeyError:
lexical_alignments = NGramAligner().align(document, summaries)
if semantic_sim_type == "Static embedding":
try:
semantic_alignments = [
example.data[f'{StaticEmbeddingAligner.__name__}:spacy:{document._.column}:spacy:{summary._.column}']
for summary in summaries
]
except KeyError:
semantic_alignments = StaticEmbeddingAligner(
semantic_sim_threshold,
semantic_sim_top_k).align(
document,
summaries
)
else:
try:
semantic_alignments = [
example.data[f'{BertscoreAligner.__name__}:spacy:{document._.column}:spacy:{summary._.column}']
for summary in summaries
]
except KeyError:
semantic_alignments = BertscoreAligner(semantic_sim_threshold,
semantic_sim_top_k).align(document,
summaries)
MainView(
document,
summaries,
semantic_alignments,
lexical_alignments,
layout,
scroll,
gray_out_stopwords,
).show(height=720)
if __name__ == "__main__":
st.set_page_config(layout="wide")
parser = argparse.ArgumentParser()
parser.add_argument('--path', type=str, default='data')
parser.add_argument('--no_clean', action='store_true', default=False,
help="Do not clean text (remove extraneous spaces, newlines).")
args = parser.parse_args()
nlp, is_lg = get_nlp()
# path = Path(args.path)
path = Path("examples/booksum/booksum_sf_sample_processed.cache")
path_dir = path.parent
all_files = set(map(os.path.basename, path_dir.glob('*')))
files = sorted([
fname for fname in all_files if not (fname.endswith(".py") or fname.startswith("."))
])
if path.is_file:
try:
file_index = files.index(path.name)
except:
raise FileNotFoundError(f"File not found: {path.name}")
else:
file_index = 0
col1, col2 = st.columns((3, 1))
filename = col1.selectbox(label="File:", options=files, index=file_index)
dataset = load_dataset(str(path_dir / filename), nlp=nlp)
dataset_size = len(dataset)
query = col2.number_input(f"Index (Size: {dataset_size}):", value=0, min_value=0, max_value=dataset_size - 1)
sidebar_placeholder_from = st.sidebar.empty()
sidebar_placeholder_to = st.sidebar.empty()
if query is not None:
example = retrieve(dataset, query, filename)
if example:
show_main(example)