import gradio as gr
from dataclasses import dataclass

from pytorch_ie.annotations import LabeledSpan
from pytorch_ie.auto import AutoPipeline
from pytorch_ie.core import AnnotationList, annotation_field
from pytorch_ie.documents import TextDocument

from spacy import displacy


@dataclass
class ExampleDocument(TextDocument):
    entities: AnnotationList[LabeledSpan] = annotation_field(target="text")


model_name_or_path = "leonhardhennig/copious_ner"

ner_pipeline = AutoPipeline.from_pretrained(model_name_or_path, device=-1, num_workers=0, taskmodule_kwargs={"include_ill_formed_predictions": False, "max_window": 512})


def predict(text):
    document = ExampleDocument(text)

    ner_pipeline(document)
    
    doc = {
        "text": document.text,
        "ents": [{
            "start": entity.start,
            "end": entity.end,
            "label": entity.label
        } for entity in sorted(document.entities.predictions, key=lambda e: e.start)],
        "title": None
    }

    html = displacy.render(doc, style="ent", page=True, manual=True, minify=True)
    html = (
        "<div style='max-width:100%; max-height:360px; overflow:auto'>"
        + html
        + "</div>"
    )
    
    return html


iface = gr.Interface(
    fn=predict,
    inputs=gr.components.Textbox(
        lines=5,
        placeholder="Enter text here or select one of the examples below.",
    ),
    outputs="html",
    title="Copious NER",
    description="NER for Taxon/Habitat/Person/GeographicalLocation trained on the COPIOUS dataset",
    examples=["OCEANIC DISPERSAL OF PLANTS. 287 fruit of the former is described as having a hard, blood-red sarcocarp, with the mesocarp traversed by stout vessels from the thin crustaceous endocarp, and this description applies fairly well to the skeleton fruit from the New Guinea drift, save that the endocarp of this is more woody. AVhen fresh this must have been nearly or quite two inches and a half in diameter. Beccari mentions, loc. cit., that he never found the fruit attached to the plant, but always picked it up from the ground. The largest fruit of the order described is that of Macrococculus pomiferus, Beccari (Malesia, i. p. 161). This is also drupaceous and globose or pyriform, and about four inches long by three inches thick. It is from the same country. BIXINE^E. Pangium edule, Beinw. Pangium edule, Reinw. ; Miq., Fl. Lid. Bat., i. 2, p. 109 ; Horsf., PI. Jav. Rar., p. 205, t. 43 ; Blunie, Rumpbia, iv. p. 20, t. 178; Blanco, Fl. Philipp., ed. 3, t. 391. New Guinea drift. Miquel states that this is generally spread in the Malayan Archipelago, from Sumatra to the Celebes, Amboina, and Ceram, though rare in Amboina ; and it is also common in the Philippines ; but hitherto it has not been found either in Continental Asia or in Australia. It is commonly cultivated by the Malays for the sake of its large, oleaginous seeds, which are an important article of food. The fruit varies considerably in size and shape, according to Blurne, who represents one nine inches long and six inches across the middle. It has a thick fleshy indehiscent pericarp, which must decay before the seeds become free. The broad, irregularly wedge-shaped seeds measure two inches or more in their greatest diameter ; and they are provided with a crustaceous testa apparently impervious to water.",
                  "Thursday Island, Torres Straits, 5-7 fms. (Coppinger)  Straits of Macassar, 11 fms. (Hinds); Damaguete, Philippine Islands Cuming for O. pupiformis) ; Philippine Islands (Cuming for C. costata); Borneo (H. Adams for O. pygmcea). The above-named and so-called species I believe to be mere variations of one and the same shell.\n\nCompare the extreme forms ( C. macrostoma and 0. pygmy), and one perceives a vast difference in outline and the number of whorls ; but even here several features in common wiu be found, namely the spotted expanded outer lip,"
                 ],
)
iface.launch()