import os 
os.system("pip install -U huggingface_hub")

from wiki_kb_qa_migrate import *

'''
zh_question = "哈利波特的血缘是什么？"
#output = from_zh_question_to_consider_queries(zh_question)
output = from_zh_question_to_consider_queries(zh_question,
                                             top_k = 32, top_p_k = 30, top_e_k = 50
                                             )
if type(output) == type((1,)):
    query_list, output = output
    zh_output = trans_output(zh_question ,output)
else:
    zh_output = None
zh_output
ranking_output(zh_question, zh_output)
'''

'''
os.system("pip uninstall httpx -y")
os.system("pip uninstall pydantic -y")
os.system("pip uninstall gradio -y")
os.system("pip install -U gradio")
os.system("pip install pydantic==2.4.2")
'''

import gradio as gr

example_sample = [
    "哈利波特的血缘是什么？",
    "哈利波特的生日是什么？",
    "占卜会议的领导者是谁？",
    "纽约卫生局的创立时间是什么？",
    "哥布林叛乱发生在什么日期？",
    "赫敏的丈夫是谁？",
]

def demo_func(zh_question):
    assert type(zh_question) == type("")
    output = from_zh_question_to_consider_queries(zh_question,
                                                 top_k = 32, top_p_k = 30, top_e_k = 50
                                                 )
    if type(output) == type((1,)):
        query_list, output = output
        zh_output = trans_output(zh_question ,output)
    else:
        zh_output = None
    #zh_output
    zh_output = ranking_output(zh_question, zh_output)
    if hasattr(zh_output, "to_dict"):
        zh_output = zh_output.head(5)
        zh_output = zh_output[["answer", "sparql_query"]].values.tolist()
    return {"output": zh_output}


demo = gr.Interface(
        fn=demo_func,
        inputs="text",
        outputs="json",
        title=f"Harry Potter Knowledge Question Answer in Chinese 🧙 demonstration",
        examples=example_sample if example_sample else None,
        cache_examples = False,
    )

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