import torch
import transformers

import gradio as gr
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
from langchain.vectorstores import Chroma
#from langchain.embeddings import OpenAIEmbeddings
#from langchain.llms import OpenAI
from langchain.chains import VectorDBQA
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings




tokenizer_nlp = AutoTokenizer.from_pretrained("FlavioBF/multi-qa-mpnet-base-dot-v1_fine_tuned_model")
model_nlp = AutoModelForQuestionAnswering.from_pretrained("FlavioBF/multi-qa-mpnet-base-dot-v1_fine_tuned_model")
#url='https://huggingface.co/FlavioBF/multi-qa-mpnet-base-dot-v1_fine_tuned_model/commit/1a9570af077d83fc8a728b0addf8a8bd276e2492'
# Load model directly
#model_name2 = "FlavioBF/multi-qa-mpnet-base-dot-v1_fine_tuned_model"
#nlp = pipeline("question-answering", model=model_nlp, tokenizer=tokenizer2)
#nlp = pipeline("fill-mask", model="FlavioBF/multi-qa-mpnet-base-dot-v1_fine_tuned_model")

persist_directory = 'https://drive.google.com/drive/folders/1jRoIBEzgT3-5Pk9eu9oWkbaUq7GxBAiV?usp=sharing'
#persist_directory = 'https://drive.google.com/file/d/1U-isHky75OdYwt0UFjYsDAh4idVyixxM/view?usp=sharing'

embedding = HuggingFaceEmbeddings()
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)
# Now we can load the persisted database from disk, and use it as normal. 
#vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)
#qa = VectorDBQA.from_chain_type(llm=OpenAI(temperature=0, model_name="gpt-3.5-turbo", openai_api_base=base_path), chain_type="stuff", vectorstore=vectordb)

context=qa.run(query)
#vectordb = Chroma.from_documents(documents=all_splits, embedding=embeddings, persist_directory="chroma_db")

retriever = vectordb.as_retriever()

qa = RetrievalQA.from_chain_type(
    llm=llm, 
    chain_type="stuff", 
    retriever=retriever, 
    verbose=True
)


#def context_rag(qa, query):
#    result = qa.run(query)
#    return result

context = context_rag(qa, query)

#def predict(context, question):
#    res = nlp({"question": question, "context": context})
#    return res["answer"], res["score"], res["Ref_mail"]


gr.Interface(
    predict,
    inputs=[
        gr.Textbox(lines=7, value=context, label="Context Paragraph"),
        gr.Textbox(lines=2, value=question, label="Question"),
    ],
    outputs=[gr.Textbox(label="Answer"), gr.Textbox(label="Score"),gr.Textbox(label="Reference mail")]
).launch(share=True)

#iface.launch(share=True)