# load required library
import os
import torch
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.llms.huggingface_pipeline import HuggingFacePipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
from langchain.document_loaders import PyPDFLoader
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate
import gradio as gr


quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16
)

model_kwargs = {'device': 'auto'}
#model_kwargs = {'device': 'cuda'}

embeddings = HuggingFaceEmbeddings(model_kwargs=model_kwargs)

#tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
#model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2", device_map='auto', quantization_config=quantization_config)
tokenizer = AutoTokenizer.from_pretrained("cmarkea/bloomz-560m-sft-chat")
model = AutoModelForCausalLM.from_pretrained("ccmarkea/bloomz-560m-sft-chat", device_map='auto', quantization_config=quantization_config)


pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150,)
llm = HuggingFacePipeline(pipeline=pipe)


pdf_link = "testrag.pdf"
loader = PyPDFLoader(pdf_link, extract_images=False)
pages = loader.load_and_split()


# Split data into chunks
text_splitter = RecursiveCharacterTextSplitter(
   chunk_size = 4000,
   chunk_overlap  = 20,
   length_function = len,
   add_start_index = True,
)
chunks = text_splitter.split_documents(pages)
print(chunks)

# Store data into database
db=Chroma.from_documents(chunks,embedding=embeddings,persist_directory="CHROMA"+pdf_link)
db.persist()

# Load the database
vectordb = Chroma(persist_directory="CHROMA"+pdf_link, embedding_function = embeddings)

# Load the retriver
retriever = vectordb.as_retriever(search_kwargs = {"k" : 3})
#qna_prompt_template="""### [INST] Instruction: You will be provided with questions and related data. Your task is to find the answers to the questions using the given data. If the data doesn't contain the answer to the question, then you must return 'Not enough information.'
qna_prompt_template="""### [INST] Instruction : tu dois répondre en francais. Des questions et des données associées vous seront fournies en francais. Votre tâche consiste à trouver les réponses aux questions en francais en utilisant les données fournies. Si les données ne contiennent pas la réponse à la question, vous devez alors renvoyer ' Pas assez d informations .'

{context}

### Question: {question} [/INST]"""

PROMPT = PromptTemplate(
   template=qna_prompt_template, input_variables=["context", "question"]
)
chain = load_qa_chain(llm, chain_type="stuff", prompt=PROMPT)


# A utility function for answer generation
def ask(question):
   context = retriever.get_relevant_documents(question)
   print(context)
   answer = (chain({"input_documents": context, "question": question}, return_only_outputs=True))['output_text']
   return answer


def summarize_pdf (input_pdf_question, custom_prompt=""):
    #print(input_pdf_question)
    summary = ask(input_pdf_question)
    return summary

#upload_file=gr.File(label="upload file", sources=['upload'],file_count="multiple")),
#    input_pdf_name = gr.Dropdown(['testrag.pdf', 'testrag.pdf'], label="file"),

gradio_app = gr.Interface(
        fn = summarize_pdf,
        inputs = gr.components.Textbox(label="Entrer la question",value="je dois repeindre un mur de 30 m2 avec de la peinture PEINTURE MATE a lille, quel est le prix final ?"),
        outputs = gr.components.Textbox(label=" "),
        title = "PDF question",
        description = "vous pouvez questionner un pdf.",
    )
demo = gradio_app.TabbedInterface([gradio_app, gradio_app], ["app1", "app2"])


if __name__ == "__main__":
    demo.launch( auth=("username", "password"), auth_message="Try this",share=True, debug=True)

#https://www.machinelearningnuggets.com/gradio-tutorial/
    
#quelle est le prix de la PLUS-VALUEGAZARGON en 2024 et en Z3
#quelle est le prix de la PLUS-VALUEGAZARGON en 2024
#quelle est le prix de la PLUS-VALUEGAZARGON en 2022
#dans quelle unité se détermine la PLUS-VALUEGAZARGON

#quel est le prix du VITRAGE TREMPE VI-TREM2D
# La question demande le prix du "VITRAGE TREMPE VI-TREM2D". Cependant, les données fournies ne contiennent pas directement cette information. Il est possible de calculer le prix en multipliant le prix par face par deux (pour avoir le prix par mètre carré) et en ajoutant les deux pourcentages de déchet (pour avoir le prix total).
#Par exemple, pour le vitrage trempe 10mm clair (VI-TREM2), le prix par face est 253 €/m² et les pourcentages de déchet sont respectivement


#quel est le prix de la peinture PEINTURE MATE sur surface endommagée ?
#quel est le prix de la peinture PEINTURE BRILLANTE sur surface endommagée en zone 1?

#quel est le prix du m2 de PEINTURE MATE sur sURFACE ENDOMMAGEE en zone 1
#quel est le prix du m2 de PEINTURE MATE sur sURFACE ENDOMMAGEE

#quel est le forfait de PEINTURE A PROJETER TYPE GOUTTELETTE  en zone 1
#quel est le forfait de PEINTURE A PROJETER TYPE GOUTTELETTE  en zone 1 pour 20 m2

#quel est le prix pour une COUCHE ANTI-MOISISSURE
#quel est le prix pour une COUCHE ANTI-MOISISSURE en zone 1
#quel est le prix pour une COUCHE ANTI-MOISISSURE en zone 2

#quel est le prix le moins cher pour une COUCHE ANTI-MOISISSURE
#quel est le premier prix  pour une COUCHE ANTI-MOISISSURE
#quel est le premier prix  pour une COUCHE ANTI-MOISISSURE appliqué a 30 m2
#quel est le premier prix  pour une COUCHE ANTI-MOISISSURE multiplié par 30 m2


#quel est le premier prix  pour une COUCHE ANTI-MOISISSURE et donne moi le résultat multiplié par 30