import os
import sys
import re

from flair.models import SequenceTagger
from flair.data import Sentence

flair_ner_model_path = "flair_model"
assert os.path.exists(flair_ner_model_path)
loaded_model: SequenceTagger = SequenceTagger.load(os.path.join(flair_ner_model_path ,"best-model.pt"))

def one_item_process(r, loaded_model):
    #assert type(r) == type(pd.Series())
    zh = r["question"]
    zh = zh.replace(" ", "").strip()
    sentence = Sentence(" ".join(list(zh)))
    loaded_model.predict(sentence)
    sentence_str = str(sentence)
    ask_spans = re.findall(r'\["(.+?)"/ASK\]', sentence_str)
    sentence = re.findall(r'Sentence: "(.+?)"', sentence_str)
    if ask_spans:
        ask_spans = ask_spans[0]
    else:
        ask_spans = ""
    if sentence:
        sentence = sentence[0]
    else:
        sentence = ""
    ask_spans, sentence = map(lambda x: x.replace(" ", "").strip(), [ask_spans, sentence])
    return ask_spans, sentence

import gradio as gr

example_sample = [
    "宁波在哪个省份？",
    "美国的通货是什么？",
]

def demo_func(question):
    assert type(question) == type("")
    ask_spans, sentence = one_item_process(
        {"question": question},
        loaded_model
    )
    return {
        "Question words": ask_spans
    }


demo = gr.Interface(
        fn=demo_func,
        inputs="text",
        outputs="json",
        title=f"Chinese Question Words extractor 🐱 demonstration",
        examples=example_sample if example_sample else None,
        cache_examples = False
    )

demo.launch(server_name=None, server_port=None)