Initial commit

app.py

@@ -0,0 +1,125 @@
+import subprocess
+try:
+    import gradio as gr
+    import spacy
+    import glirel
+except:
+    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"])
+
+try:
+    import spacy
+    spacy.load("en_core_web_sm")
+except:
+    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(server_port=9989)