# %% auto 0
__all__ = ['text_splitter', 'embeddings', 'hub_llm', 'pdf_example_1', 'pdf_example_2', 'question_example_1', 'question_example_2',
           'title', 'description', 'file_upload', 'question', 'output', 'split_pdf', 'ask_pdf']

import dotenv
from dotenv import load_dotenv
load_dotenv()

# Import qa chain
from langchain.chains.question_answering import load_qa_chain

# Import gradio for UI/app interface
import gradio as gr

# Import PDF Loaders
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

# Create an instance of RecursiveCharacterTextSplitter with your desired chunk_size and chunk_overlap
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=250)

def split_pdf(pdf):
  if pdf is not None:
      # load pdf
    loader = PyPDFLoader(pdf)
      # Split pages from pdf (SEE: https://api.python.langchain.com/en/latest/document_loaders/langchain.document_loaders.pdf.PyPDFLoader.html)
    chunks = loader.load_and_split(text_splitter=text_splitter)
    return chunks

# Import HuggingFace as main LLM service
from langchain.embeddings import HuggingFaceEmbeddings

embeddings = HuggingFaceEmbeddings(
    model_name = "bert-base-nli-mean-tokens",
    model_kwargs = {'device': 'cpu'},
    encode_kwargs = {'normalize_embeddings': True}
)

# Import FAISS as the vector store
from langchain.vectorstores import FAISS

# import transformers
# from transformers import AutoTokenizer, AutoModelForCausalLM
from langchain import HuggingFaceHub
import torch

hub_llm = HuggingFaceHub(
    repo_id = 'tiiuae/falcon-7b-instruct',
    # repo_id = 'tiiuae/falcon-7b-instruct',
    model_kwargs = {'temperature' : 0.01,  # 'randomness' of outputs, 0.0 is the min and 1.0 the max
                    # 'top_p': 0.15,  # select from top tokens whose probability add up to x%
                    'top_k': 25,  # select from top x tokens
                    'max_new_tokens': 300,  # mex number of tokens to generate in the output
                    'repetition_penalty': 999  # without this output begins repeating
                    }
)

chain = load_qa_chain(hub_llm, chain_type="stuff")

# App Function
def ask_pdf(pdf, user_question):
  if pdf is not None and user_question != '':
      chunks = split_pdf(pdf) # split pdf into smaller chunks
      store = FAISS.from_documents(chunks, embeddings) # Load documents into vector database
      docs = store.similarity_search(user_question) # take user input and look for chunks that might contain the answer
    # run chain
      response = chain.run(input_documents=docs, question=user_question)

  return response

pdf_example_1 = '/attention.pdf' 
pdf_example_2 = '/bert.pdf'
question_example_1 = "What does 'attention' mean in this context?"
question_example_2 = "What does BERT stand for?"

title="LangChain project: Conversational Retrieval Chain"
description="Upload a PDF file and ask a question related to its content."

# Create Gradio interface
file_upload = gr.File(file_types = ['.pdf'],
                              file_count='single',
                              label="Upload a PDF file")
question = gr.Textbox(label="Write your question here:", show_copy_button= True)
output = gr.Textbox(label="Your answer:")

gr.Interface(
    fn=ask_pdf,
    inputs=[file_upload, question],
    outputs=output,
    title=title,
    description=description,
    examples=[[pdf_example_1, question_example_1], [pdf_example_2, question_example_2]]
).launch() #share=True