import arxiv
import gradio as gr
from langchain.document_loaders import OnlinePDFLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.llms import HuggingFaceHub
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA

repo_id = 'mistralai/Mistral-7B-v0.1'
client = arxiv.Client()

def loading_paper(): return 'Loading...'

def paper_changes(paper_id):
    paper = next(arxiv.Client().results(arxiv.Search(id_list=[paper_id])))
    loader = OnlinePDFLoader(paper.download_pdf())
    documents = loader.load()
    text_splitter = CharacterTextSplitter(chunk_size=1024, chunk_overlap=0)
    texts = text_splitter.split_documents(documents)
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
    db = Chroma.from_documents(texts, embeddings, persist_directory="chroma_db")
    retriever = db.as_retriever()
    llm = HuggingFaceHub(repo_id=repo_id, model_kwargs={'temperature': 0.5, 'max_new_tokens': 2096})
    global qa
    qa = RetrievalQA.from_chain_type(llm=llm, chain_type='stuff', retriever=retriever, return_source_documents=True)
    return 'Ready!!'

def add_text(history, text):
    history = history + [(text, None)]
    return history, ''

def bot(history):
    response = infer(history[-1][0])
    history[-1][1] = response['result']
    return history

def infer(question):
    result = qa({'query': question})
    return result

with gr.Blocks(theme='WeixuanYuan/Soft_dark') as demo:
    with gr.Column():
        chatbot = gr.Chatbot([], elem_id='chatbot')

        with gr.Row():
            paper_id = gr.Textbox(label='ArXiv Paper Id', placeholder='1706.03762')
            langchain_status = gr.Textbox(label='Status', placeholder='', interactive=False)
            load_paper = gr.Button('Load Paper to LangChain')

        with gr.Row():
            question = gr.Textbox(label='Question', placeholder='Type your query...')
            submit_btn = gr.Button('Submit')

    load_paper.click(paper_changes, inputs=[paper_id], outputs=[langchain_status], queue=False)
    question.submit(add_text, [chatbot, question], [chatbot, question]).then(bot, chatbot, chatbot)
    submit_btn.click(add_text, [chatbot, question], [chatbot, question]).then(bot, chatbot, chatbot)

demo.launch(share=True)