import gradio as gr

from pathlib import Path
from supar import Parser
from spacy import displacy
from spacy.tokens import Doc, Span
import spacy

proj_dir = Path(__file__).parent
model_choices = sorted([str(model.name) for model in (proj_dir / 'models').glob('*')])


def sentence_diagram(model_name, text, progress=gr.Progress(track_tqdm=True)):
    parser = Parser.load(f'./models/{model_name}')

    Span.set_extension("con_tree", getter=lambda x: parser.predict([i.text for i in x], verbose=False)[0], force=True)
    nlp = spacy.load('en_core_web_sm')
    doc = nlp(text)

    svg = displacy.render(doc, style="dep")
    output_path = Path("sentence.svg")
    output_path.open("w", encoding="utf-8").write(svg)
    return output_path


with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("""
        # Purpose
        Way back in 7th grade, my english teacher **Brother Hill** would always disclaim our sentence diagram lessons with: 
        "*you probably wont be doing these in 20 years*". A few of us being middle schoolers would love to contradict this. 
        Unfortunately he passed away in 2015, so I thought this would be a nice tribute.
        
        # Instructions
        1. Choose a model:
            - `ptb.biaffine.dep.roberta` is slower but marginally better
            - `ptb.biaffine.dep.lstm.char` is faster but marginally worse
        2. Write your sentence
        3. Click Run!
        """)

        gr.HTML("<img src='file=https://huggingface.co/spaces/derek-thomas/sentence_diagrams/resolve/main/media/bro_hill.jpg' />") # work

    model_name = gr.Dropdown(choices=model_choices, label='Model', value=model_choices[0])
    text_in = gr.Textbox(label='Sentence(s) to diagram',
                         value='You were a great teacher, and Im thankful for the impact you had in my life!')
    button = gr.Button('Run!')
    html_out = gr.Image()
    gr.Markdown("""
    # Information
    ##### This doesnt look like the sentences we used to do!
    
    There are some slight differences between 
    [Reed-Kellogg](https://blog.ung.edu/press/classroom-grammar-an-introduction-to-the-reed-kellogg-system/)
    and [Dependency Parsing](https://en.wikipedia.org/wiki/Dependency_grammar)
    in both presentation and linquistic analysis as shown [here](https://en.wikipedia.org/wiki/Sentence_diagram),
    but they are similar enough for me not to mind too much. 
    
    ##### How did you do this?
    
    I chose a state of the art **Dependency Parsing** [model](https://github.com/yzhangcs/parser) as of ~2 years ago. 
    I believe this has been [surpassed](https://paperswithcode.com/sota/dependency-parsing-on-penn-treebank) 
    in recent years. 
    
    Dependency Parsing was a popular task in NLP to feed to models to improve performance, but in the age of the 
    [transformer](https://arxiv.org/abs/1706.03762) it's rarelu used in anymore.
    
    Then I deployed this in a [Gradio App](https://gradio.app) on a [Hugging Face Space](https://huggingface.co/spaces).
    
    # To Brother Hlll
    Thanks for being a great teacher. As an adult I appreciate even more that you invested in so many of us,
    yet you didnt get to witness a lot of the results. 
    
    > One generation plants the trees, and another gets the shade.
    >
    > ~ [Chinese Proverb](https://rotarycluboflahainasunset.org/stories/one-generation-plants-the-trees-and-another-gets-the-shade-(chinese-proverb))
    
    I have a lot of fond memories of you from PE, English, and Home Repair, and I wish we
    could have connected before you passed away. 
    
    Thanks again,
    
    Derek
    """)

    button.click(sentence_diagram,
                 inputs=[model_name, text_in],
                 outputs=html_out)

if __name__ == '__main__':
    demo.queue().launch(show_error=True)