import gradio as gr
import BertQuery


def queryColelctionByText(text):
    docs = BertQuery.queryCollection(BertQuery.text_to_vector(text))
    if docs is None or len(docs) == 0:
        return '查询失败'
    # docs 返回id和fields
    result = []
    result.append("videoId\t标题\t回流人数\t分享次数\t分享人数\t曝光次数\t曝光人数\t播放次数\t播放人数\t相似度")
    for doc in docs:
        videoId = doc.id
        title = doc.fields['title']
        rntHeadCount = doc.fields['rntHeadCount']
        shareCount = doc.fields['shareCount']
        shareHeadCount = doc.fields['shareHeadCount']
        exposureCount = doc.fields['exposureCount']
        exposureHeadCount = doc.fields['exposureHeadCount']
        playCount = doc.fields['playCount']
        playHeadCount = doc.fields['playHeadCount']
        result.append(
            f'{videoId}\t{title}\t{rntHeadCount}\t{shareCount}\t{shareHeadCount}\t{exposureCount}\t{exposureHeadCount}\t{playCount}\t{playHeadCount}\t{doc.score}')
    return '\n\n'.join(result)


iface = gr.Interface(fn=queryColelctionByText,
                     inputs=gr.components.Textbox(
                         lines=7, label="请输入文本内容"),
                     outputs="text",
                     title="视频内容向量检索，相似度匹配")
iface.launch(share=True)