import gradio as gr
import spacy 
from spacy import displacy
import pandas as pd 

df = pd.read_json("cliff_gradio.jsonl", orient="index")
colors = {"ext": "#aa9cfc", "int": "#ff6961", "wor": "#7aecec"}
options = {"colors": colors}

nlp = spacy.blank("en")

def viz(topic_id):
    row = df.loc[topic_id]
    summary = row["summary"]    
    labels = row["labels"]
    score = row["score"]
    ex = [{"text": summary,
       "ents": [{"start": token.idx, "end": token.idx + token.__len__() + 1, "label": label[:3]} for token, label in zip(nlp(summary), labels) if label != "correct"],
       "title": None}]
    return row["article"], displacy.render(ex, style="ent", manual=True, options=options), score


    

demo = gr.Interface(fn=viz,
                    inputs=[gr.Dropdown(df.index.tolist(), label="Select Topic Id")],
                    outputs=[gr.Text(label="Source article"), 
                             gr.HTML(label="summary"), 
                             gr.Number(label="score")])



demo.launch()