Huu076/InformationExtraction

1from typing import Dict, Union2from gliner import GLiNER3import gradio as gr4 5model = GLiNER.from_pretrained("knowledgator/gliner-multitask-large-v0.5").to('cpu')6 7 8def merge_entities(entities):9    if not entities:10        return []11    merged = []12    current = entities[0]13    for next_entity in entities[1:]:14        if next_entity['entity'] == current['entity'] and (next_entity['start'] == current['end'] + 1 or next_entity['start'] == current['end']):15            current['word'] += ' ' + next_entity['word']16            current['end'] = next_entity['end']17        else:18            merged.append(current)19            current = next_entity20    merged.append(current)21    return merged22 23def process(24    prompt:str, text, threshold: float, nested_ner: bool, labels: str = ["match"]25) -> Dict[str, Union[str, int, float]]:26    text = prompt + "\n" + text27    r = {28        "text": text,29        "entities": [30            {31                "entity": entity["label"],32                "word": entity["text"],33                "start": entity["start"],34                "end": entity["end"],35                "score": 0,36            }37            for entity in model.predict_entities(38                text, labels, flat_ner=not nested_ner, threshold=threshold39            )40        ],41    }42    r["entities"] =  merge_entities(r["entities"])43    return r44 45with gr.Blocks(title="Open Information Extracting") as open_ie_interface:46    prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")47    input_text = gr.Textbox(label="Text input", placeholder="Enter your text here")48    threshold = gr.Slider(0, 1, value=0.3, step=0.01, label="Threshold", info="Lower the threshold to increase how many entities get predicted.")49    nested_ner = gr.Checkbox(label="Nested NER", info="Allow for nested NER?")50    output = gr.HighlightedText(label="Predicted Entities")51    submit_btn = gr.Button("Submit")52    53    theme=gr.themes.Base()54 55    input_text.submit(fn=process, inputs=[prompt, input_text, threshold, nested_ner], outputs=output)56    prompt.submit(fn=process, inputs=[prompt, input_text, threshold, nested_ner], outputs=output)57    threshold.release(fn=process, inputs=[prompt, input_text, threshold, nested_ner], outputs=output)58    submit_btn.click(fn=process, inputs=[prompt, input_text, threshold, nested_ner], outputs=output)59    nested_ner.change(fn=process, inputs=[prompt, input_text, threshold, nested_ner], outputs=output)60 61 62if __name__ == "__main__":63    64    open_ie_interface.launch()