import math, torch, gradio as gr

from lex_rank import LexRank
from lex_rank_new_model import LexRankNewModel
from sentence_transformers import SentenceTransformer, util


# ---===--- instances ---===---
embedder = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')
lex = LexRank()
lex_new_model = LexRankNewModel()


# 摘要方法
def extract_handler(content):
    summary_length = math.ceil(len(content) / 10)
    sentences = lex.find_central(content, num=summary_length)
    output = ""
    for index, sentence in enumerate(sentences):
        output += f"{index}: {sentence}\n"
    return output


# 摘要方法
def extract_handler_new_model(content):
    summary_length = math.ceil(len(content) / 10)
    sentences = lex_new_model.find_central(content, num=summary_length)
    output = ""
    for index, sentence in enumerate(sentences):
        output += f"{index}: {sentence}\n"
    return output


# 相似度检测方法
def similarity_search(queries, doc):
    doc_list = doc.split('\n')
    query_list = queries.split('\n')

    corpus_embeddings = embedder.encode(doc_list, convert_to_tensor=True)
    top_k = min(5, len(doc_list))
    output = ""
    for query in query_list:
        query_embedding = embedder.encode(query, convert_to_tensor=True)
        # We use cosine-similarity and torch.topk to find the highest 5 scores
        cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0]
        top_results = torch.topk(cos_scores, k=top_k)
        output += "\n\n======================\n\n"
        output += f"Query: {query}"
        output += "\nTop 5 most similar sentences in corpus:\n"
        for score, idx in zip(top_results[0], top_results[1]):
            output += f"{doc_list[idx]}(Score: {score})\n"
    return output


#  web ui
with gr.Blocks() as app:
    gr.Markdown("从下面的标签选择测试模块 [摘要生成,相似度检测]")
    with gr.Tab("LexRank-mpnet"):
        text_input_1 = gr.Textbox(label="请输入长文本:", lines=10, max_lines=1000)
        text_button_1 = gr.Button("生成摘要")
        text_output_1 = gr.Textbox(label="摘要文本(长度设置为原文长度的1/10)", lines=10)
    with gr.Tab("LexRank-distiluse"):
        text_input_2 = gr.Textbox(label="请输入长文本:", lines=10, max_lines=1000)
        text_button_2 = gr.Button("生成摘要")
        text_output_2 = gr.Textbox(label="摘要文本(长度设置为原文长度的1/10)", lines=10)
    with gr.Tab("相似度检测"):
        with gr.Row():
            text_input_query = gr.Textbox(lines=10, label="查询文本")
            text_input_doc = gr.Textbox(lines=20, label="逐行输入待比较的文本列表")
        text_button_similarity = gr.Button("对比相似度")
        text_output_similarity = gr.Textbox()

    text_button_1.click(extract_handler, inputs=text_input_1, outputs=text_output_1)
    text_button_2.click(extract_handler_new_model, inputs=text_input_2, outputs=text_output_2)
    text_button_similarity.click(similarity_search, inputs=[text_input_query, text_input_doc], outputs=text_output_similarity)

app.launch(
    # enable share will generate a temporary public link.
    # share=True,
    # debug=True,
    # auth=("qee", "world"),
    # auth_message="请登陆"
           )