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})


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.inputs.Textbox(
        lines=5,
        default="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.\nPangium 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. Only one seed is in the collection, and that is sound, though it has probably lost its germinative power. In its present state it floats with about half its body submerged. Good flowering specimens of this tree are wanted for the Kew Herbarium. \nDIPTEROCARPE.E. Vateria papuana, Dyer. (Plate LXIV., B.) Vateria papuana, Dyer, ante, p. 123. Vatica papuana, Dyer in Journ. Bot., 1878, p. 100. Vateria sp., Beccari in D'Albertis, New Guinea, ii. p. 394. \nFruits from the New Guinea drift, and sea-beach Arrou Islands ; flowering specimens and fruits from Bauioi, Southern New Guinea. In the first place, only flowering specimens of this very distinct Dipterocarp were communicated to Mr Dyer, and as the floral structure agreed with Vatica, it was referred",
    ),
    outputs="html",
)
iface.launch()