import pdfplumber
import re
import numpy as np
import tensorflow_hub as hub
import openai
import os
import tensorflow_text
from sklearn.neighbors import NearestNeighbors
import gradio as gr
import requests
import json


#这里填写调用openai需要的密钥
openai.api_key = '9481961416fa4c8e883047c5679cf971'
openai.api_base = 'https://demopro-oai-we2.openai.azure.com/' 
openai.api_type = 'azure'
openai.api_version = '2022-12-01'

#将嵌套的列表展平
def flatten(_2d_list):
    flat_list = []
    for element in _2d_list:
        if type(element) is list:
            for item in element:
                flat_list.append(item)
        else:
            flat_list.append(element)
    return flat_list


#将pdf文档按段落分
def pdf_to_chunks(path,ver):
    doc = pdfplumber.open(path)
    pages = doc.pages
    #删掉目录和封面
    del pages[0];del pages[0];del pages[0];del pages[0]; 
    text_list=[]
    for page,d in enumerate(pages):
        d=d.extract_text()
        #正则表达式，以 标点符号+换行 为分段点 \n表示换行符
        d=re.split('[。！？!?.]\n',d)
        for i in range(len(d)):
            d[i]=d[i].replace('\n','')
            #在每一段的开头加上 [七年级上册10页] 这样的标记，让chatgpt之后能够识别
            d[i]=f'[七年级{ver}册{page}页]'+d[i]
        text_list.append(d)
        
    doc.close()   
    #上面的操作完毕后，得到的text_list是二维列表 [[页码1的段落],[页码2的段落],...,[页码n的段落]] 
    #因为我们已经在每段开头标记好了页码，所以不需要二重列表结构来表示页码，为了之后方便操作，需要展平为[段落1,段落2,...,段落n]
    return flatten(text_list)

history1=pdf_to_chunks('人教部编版历史7年级上册.pdf','上')
history2=pdf_to_chunks('人教部编版历史7年级下册.pdf','下')
history=history1+history2

#定义语义搜索类
class SemanticSearch:
    
    def __init__(self):
        #类初始化，使用google公司的多语言语句编码，第一次运行时需要十几分钟的时间下载
        self.use = hub.load('https://tfhub.dev/google/universal-sentence-encoder-multilingual/3')
        self.fitted = False
    
    #K近邻算法，找到与问题最相似的 k 个段落,这里的 k 即n_neighbors=10
    def fit(self, data,  n_neighbors=10):
        self.data = data
        #self.embeddings即整本书的段落转化成的向量
        self.embeddings = self.use(data)
        n_neighbors = min(n_neighbors, len(self.embeddings))
        self.nn = NearestNeighbors(n_neighbors=n_neighbors)
        self.nn.fit(self.embeddings)
        self.fitted = True
    
    #定义了该方法后，实例就可以被当作函数调用,text参数即用户提出的问题，inp_emb为其转化成的向量
    def __call__(self, text, return_data=True):
        inp_emb = self.use([text])
        
        
        neighbors = self.nn.kneighbors(inp_emb, return_distance=False)[0]
        
        if return_data:
            return [self.data[i] for i in neighbors]
        else:
            return neighbors
    
    
#openai的api接口，engine参数为我们选择的大语言模型，prompt即提示词
def generate_text(prompt, engine="text-davinci-003"):
    completions = openai.Completion.create(
        engine=engine,
        prompt=prompt,
        max_tokens=512,
        n=1,
        stop=None,
        temperature=0.7,
    )
    message = completions.choices[0].text
    return message


def generate_answer(question):
    #匹配与问题最相近的n个段落，前面定义了n=10
    topn_chunks = recommender(question)
    prompt = ""
    prompt += 'search results:\n\n'
    
    #把匹配到的段落加进提示词
    for c in topn_chunks:
        prompt += c + '\n\n'
    
    #提示词
    prompt += '''
    Instructions: 如果搜索结果中找不到相关信息，只需要回答'未在该文档中找到相关信息'。
    如果找到了相关信息，请使用中文回答，回答尽量精确简洁。并在句子的末尾使用[七年级上册/七年级下册页码]符号引用每个参考文献（每个结果的开头都有这个编号)
    如果不确定答案是否正确，就仅给出相似段落的来源，不要回复错误的答案。
    \n\nQuery: {question}\nAnswer: 
    '''
    
    prompt += f"Query: {question}\nAnswer:"
    answer = generate_text(prompt,"text-davinci-003")
    return answer


recommender = SemanticSearch()
recommender.fit(history)

#以下为web客户端搭建，运行后产生客户端界面
def ask_api(question):
    
    if question.strip() == '':
        return '[ERROR]: 未输入问题'

    return generate_answer(question)

title = 'ChatHistory'
description = """ 你好，我是ChatHistory，我刚刚成长到七年级，可以回答七年级历史书中的问题了，快来问问我吧
"""

with gr.Blocks() as demo:
    gr.Markdown(f'<center><h1>{title}</h1></center>')
    gr.Markdown(description)

    with gr.Row():
        with gr.Group():
            question = gr.Textbox(label='请输入你的问题')
            btn = gr.Button(value='提交')
            btn.style(full_width=True)

        with gr.Group():
            answer = gr.Textbox(label='ChatHistory给你的回答是:')

        btn.click(
            ask_api,
            inputs=[question],
            outputs=[answer]
        )

#参数share=True会产生一个公开网页，别人可以通过访问该网页使用你的模型，前提是你需要正在运行这段代码（将自己的电脑当作服务器）
demo.launch()