import argparse

import gradio as gr
from search_online import OnlineSearcher

K = 10

# default input
DEFAULT_QUERY_MRTYDI="""사용자 질의와 관련된 문단들을 Wikipedia 한국어 코퍼스에서 반환합니다.\n
예를들어\n
- 최초로 전기 자동차를 개발한 기업은 어디야?
- 스쿠버 다이빙 잠수 가능 깊이는 최대 몇미터인가요?
- 독점규제법 상의 기업결합심사기준은 기업결합을 어떻게 구분하여 경쟁제한성을 판단하나요?
"""


# manual arguments (FIXME)
args = argparse.Namespace
args.index_type='hybrid'
args.index="/root/indexes/mrtydi-korean/sparse,/root/indexes/mrtydi-korean/dense"
args.encoder="castorini/mdpr-question-nq"
args.device="cuda:0"
args.alpha=1000000
args.normalization=True
args.lang_abbr='ko'

# initialize qabot
print("initialize Mr.tydi retrieval bot")
searcher = OnlineSearcher(args)

def Retrieve(query):
    hits = searcher.search(query, K)

    result = searcher.print_result(hits, K)
    
    return result

gr.Interface(
    fn=Retrieve, 
    inputs=[
            gr.inputs.Textbox(
            default=DEFAULT_QUERY_MRTYDI, 
            label="질의"),
        ],
    outputs=[
            gr.inputs.Textbox(
            default="검색된 문단들과 점수를 출력합니다", 
            label="검색 결과")
        ], 
    title="Hybrid (term + neural) 검색",
    theme='dark-grass',
    description=f"특정 지식베이스에 대해 학습한 검색 시스템을 테스트합니다.\n지식베이스에 맞추어 재학습이 가능하며, 현재 데모에서는 Wikipedia 한국어 코퍼스 (2020 dump, mr.Tydi ver, #문단=1496126)에 대해 학습한 모델을 사용합니다. \n\n neural: castorini/mdpr-[passage,question]-nq, term: BM25"
).launch(share=True)