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 HuggingFaceHubEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
import os

#os.environ["HUGGINGFACEHUB_API_TOKEN"] = ""


def file_upload_click(pdf_doc):  
    loader = OnlinePDFLoader(pdf_doc.name)
    documents = loader.load()
    text_splitter = CharacterTextSplitter(chunk_size=300, chunk_overlap=0)
    texts = text_splitter.split_documents(documents)
    embeddings = HuggingFaceHubEmbeddings()
    db = Chroma.from_documents(texts, embeddings)
    retriever = db.as_retriever()
    llm = HuggingFaceHub(repo_id="OpenAssistant/oasst-sft-1-pythia-12b", model_kwargs={"temperature":0.1, "max_new_tokens":250})
    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):
    query=history[-1][0]
    response = qa({"query": query})
    history[-1][1] = response['result']
    return history


with gr.Blocks() as demo:
    status_label = gr.Label(value='Start')
    file_upload = gr.File(label="Uplaod pdf", file_types=['.pdf'], type="file")
    file_upload_button= gr.Button('upload file')
    chatbot = gr.Chatbot([], elem_id="chatbot").style(height=350)
    question = gr.Textbox(label="Question", placeholder="Type your question and click submit")
    submit_btn = gr.Button("Send message")
    file_upload_button.click(file_upload_click, inputs=[file_upload], outputs=[status_label], queue=False)
    submit_btn.click(add_text, [chatbot, question], [chatbot, question], queue=False).then(
        bot, chatbot, chatbot
    )


demo.queue()
demo.launch()