import subprocess
subprocess.run(["pip", "uninstall", "gradio"])
subprocess.run(["pip", "install", "gliner"])
subprocess.run(["pip", "install", "gradio==4.31.5"])

import gradio as gr
from typing import Dict, Union
from gliner import GLiNER
import gradio as gr

model = GLiNER.from_pretrained("BioMike/logical-gliner-large").to("cpu")

qa_examples = [
    [
        "",
        "For a student to graduate, they must complete all their required courses and pass the final exam. John has completed all his required courses but failed the final exam. Answer options: 1. John can graduate 2. John cannot graduate 3. John completed all his required courses 4. John passed the final exam",
        0.5,
        False
    ],
    [
        "",
        "(P ∨ Q) → R, (R ∧ S) → T, ¬T, P. Answer options: 1. ¬R 2. ¬S 3. ¬Q 4. R 5. S 6. T",
        0.5,
        False
    ],
    [
        "",
        "(A → B), (B → (C ∧ D)), ¬C, E → F, A, ¬F. Answer options: 1. ¬A 2. ¬B 3. ¬E 4. F",
        0.5,
        False
    ]]

def merge_entities(entities):
    if not entities:
        return []
    merged = []
    current = entities[0]
    for next_entity in entities[1:]:
        if next_entity['entity'] == current['entity'] and (next_entity['start'] == current['end'] + 1 or next_entity['start'] == current['end']):
            current['word'] += ' ' + next_entity['word']
            current['end'] = next_entity['end']
        else:
            merged.append(current)
            current = next_entity
    merged.append(current)
    return merged

def process(
    question:str, text, threshold: float, nested_ner: bool, labels: str = ["answer"]
) -> Dict[str, Union[str, int, float]]:
    text = question + "\n" + text
    r = {
        "text": text,
        "entities": [
            {
                "entity": entity["label"],
                "word": entity["text"],
                "start": entity["start"],
                "end": entity["end"],
                "score": 0,
            }
            for entity in model.predict_entities(
                text, labels, flat_ner=not nested_ner, threshold=threshold
            )
        ],
    }
    r["entities"] =  merge_entities(r["entities"])
    return r

with gr.Blocks(title="Question Answering Task") as qa_interface:
    question = gr.Textbox(label="Question", placeholder="Enter your question here")
    input_text = gr.Textbox(label="Text input", placeholder="Enter your text here")
    threshold = gr.Slider(0, 1, value=0.3, step=0.01, label="Threshold", info="Lower the threshold to increase how many entities get predicted.")
    nested_ner = gr.Checkbox(label="Nested NER", info="Allow for nested NER?")
    output = gr.HighlightedText(label="Predicted Entities")
    submit_btn = gr.Button("Submit")
    examples = gr.Examples(
    qa_examples,
    fn=process,
    inputs=[question, input_text, threshold, nested_ner],
    outputs=output,
    cache_examples=True
    )
    theme=gr.themes.Base()

    input_text.submit(fn=process, inputs=[question, input_text, threshold, nested_ner], outputs=output)
    question.submit(fn=process, inputs=[question, input_text, threshold, nested_ner], outputs=output)
    threshold.release(fn=process, inputs=[question, input_text, threshold, nested_ner], outputs=output)
    submit_btn.click(fn=process, inputs=[question, input_text, threshold, nested_ner], outputs=output)
    nested_ner.change(fn=process, inputs=[question, input_text, threshold, nested_ner], outputs=output)

qa_interface.queue()
qa_interface.launch(debug=True, share=True)