import subprocess

subprocess.run(["pip", "install", "gradio==4.31.5"])
subprocess.run(["pip", "install", "spacy"])
subprocess.run(["pip", "install", "glirel"])
subprocess.run(["pip", "install", "scipy==1.10.1"])
subprocess.run(["pip", "install", "numpy==1.26.4"])

subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
subprocess.run(["python", "-m", "spacy", "download", "en_core_web_md"])
subprocess.run(["python", "-m", "spacy", "download", "en_core_web_lg"])




from typing import Dict, Union
import gradio as gr
from glirel import GLiREL
import spacy

examples = [
    [
        "Amazon, founded by Jeff Bezos, is a leader in e-commerce and cloud computing. The company has also ventured into artificial intelligence and digital streaming.",
        "en_core_web_sm",
        "Founded_By, Located_In, Produces, Operates_In, Works_With, Known_For, Headquartered_In, Partnership_With, Innovates_In, Established_In",
    ],
    [
        "J.K. Rowling, the author of the Harry Potter series, has significantly impacted modern literature. Her books have been translated into numerous languages and adapted into successful films.",
        "en_core_web_sm",
        "Translated_Into, Adapted_Into, Born_In, Author_Of, Known_For, Works_With, Located_In, Writes_For, Produced_By, Published_By"
    ],
    [
        "Apple Inc. was founded by Steve Jobs, Steve Wozniak, and Ronald Wayne in April 1976. The company is headquartered in Cupertino, California.",
        "en_core_web_sm",
        "CO_FOUNDER, HEADQUARTERED_IN, FOUNDED_BY, LOCATED_IN, ESTABLISHED_IN, PARTNERSHIP_WITH, WORKS_WITH, KNOWN_FOR"
    ]

]


# Load the relation extraction model
rel_model = GLiREL.from_pretrained("jackboyla/glirel_beta")

# Function to perform Named Entity Recognition
def perform_ner(text, model_name):
    nlp = spacy.load(model_name)
    doc = nlp(text)
    return doc

# Function to extract relations
def extract_relations(tokens, ner, labels):
    relations = rel_model.predict_relations(tokens, labels, threshold=0.0, ner=ner, top_k=1)
    sorted_data_desc = sorted(relations, key=lambda x: x['score'], reverse=True)
    return sorted_data_desc

def format_ner(text, ner):
    if isinstance(ner[0], spacy.tokens.Span):
        # if ner is spacy entities; otherwise we assume the format is correct
        ner = [[ent.start_char, ent.end_char, ent.label_, ent.text] for ent in ner]
    return {
        "text": text,
        "entities": [
            {
                "entity": entity[2],
                "word": entity[3],
                "start": entity[0],
                "end": entity[1],
                "score": 0,
            }
            for entity in ner
        ],
    }

# Gradio Interface
def process(text, model_name, labels):
    doc = perform_ner(text, model_name)
    tokens = [token.text for token in doc]
    ner = [[ent.start, (ent.end-1), ent.label_, ent.text] for ent in doc.ents]
    labels = labels.split(',')
    relations = extract_relations(tokens, ner, labels)
    print(relations)
    formatted_ner = format_ner(doc.text, doc.ents)
    formatted_rel = ""
    for item in relations:
        formatted_rel += f"{item['head_text']} --> {item['label']} --> {item['tail_text']} \t\t| score: {item['score']}\n" 
    return formatted_ner, formatted_rel

# Gradio App Layout
with gr.Blocks() as demo:

    gr.Markdown("# 🕵️‍♀️GLiREL: Zero-Shot Relation Extraction")
    gr.Markdown("GitHub: https://github.com/jackboyla/GLiREL")

    text_input = gr.Textbox(label="Input Text", value="Jack lives in London but he was born in Mongolia.")
    model_name_input = gr.Dropdown(choices=["en_core_web_sm", "en_core_web_md", "en_core_web_lg"], label="NER Model", value="en_core_web_sm")
    labels_input = gr.Textbox(label="Relation Labels (comma-separated)", value="country of origin, licensed to broadcast to, father, followed by, characters")

    ner_output = gr.HighlightedText(label="NER")
    rel_output = gr.Textbox(label="Relation Extraction Results")

    extract_button = gr.Button("Extract Relations")
    extract_button.click(
        fn=process,
        inputs=[text_input, model_name_input, labels_input],
        outputs=[ner_output, rel_output]
    )

    examples = gr.Examples(
        examples,
        fn=process,
        inputs=[text_input, model_name_input, labels_input],
        outputs=[ner_output, rel_output],
        cache_examples=True,
    )


if __name__ == "__main__":
    demo.launch()