import gradio as gr

#|------------------------
#| requirements.txt file
#|------------------------
#| torch
#| transformers

import torch
import pandas as pd
import transformers
import os
import gradio as gr
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain_openai import OpenAI
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain_openai import OpenAI
from langchain.llms import OpenAI
from transformers import AutoTokenizer,pipeline
from langchain.chat_models import ChatOpenAI
from langchain.vectorstores import Chroma
from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings

#HF_TOKEN = os.getenv('HF_TOKEN')
#OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')

api_key = os.getenv("OPEN_API_KEY")
client = OpenAI(api_key=api_key)

title = "QUESTION ANSWERIGN WITH RAG TECHNIQUES"
description = """
This Space uses:\n
- FlavioBF/multi-qa-mpnet-base-dot-v1_fine_tuned_model \n
- ChatGPT 3.5-turbo \n
to get insights from the Enron Mail dataset. In particular:\n
The Question Asnswer model "multi-qa-mpnet-base-dot-v1" has been been used to create embeddings of a small subset of the dataset and, then, trained over a sampled dataset of aapx. 500k istances of the original Enron dataset 
Embedded content (stored in Chroma DB) is used to retrieve context that will be used by the downnstream processing using similarity analysis (metric for similarity = dot product). \n
NOTE: The chunk size of 500 chars used in the text splitters, is probably too small to capture properly sentences in an effective way, neverthless it has been kept. \n

Finally, to answer to questions from Enron dataset, both models are using the context generated using RAG technique (retriever). \n

REMARK: due to the limited storage capacity the context can be generated only over a limited number of mails.\n 
The GPT 3.5 turbo model has been instructed to avoid to make up answers in case contecxt is not clear
"""

examples=[
	['Has Energy Derivative risk management  been discussed?'],
    ['Has how Adjustment Payment works in combination with inc dec pricing been discussed?'],
    ['Would have Enron stockholders wanted that Enron participates in discussions between the PX and the Chairman of FERC?'],
    ['How ENA could affect financial liability'],
    ['Next action after the first round of RTO filings and responses complete'],
    ['Who has written the draft of a GISB for review?'],
    ['When it has been discussed the gas and power price situation in California'],
    ['were the data provided to Arthur Andersen for model validation request sufficient?'],
	['Tell me about Project Stanley is a unique, sensitive project that is being handled'],
	['Who has drafted Here is AReM response to the 10-19 Angelides letter?'],   
    ]

#query2='how Top Gas Gorillas is performing?'
#context2='Note the discussion on market share halfway through the article. Top Gas Gorillas Show Strong Volume Growth The year 2000 was a banner year for the top players in gas marketing, with huge increases in gas prices, enormous volatility, continuing growth in sales volumes and major potential for profits.'

model_name = 'sentence-transformers/multi-qa-mpnet-base-dot-v1'

if torch.cuda.is_available():
    model_kwargs = {"device": "cuda"}
else:
    model_kwargs = {"device": "cpu"}

embeddings = HuggingFaceEmbeddings(model_name=model_name, model_kwargs=model_kwargs)
data_emb=pd.read_csv('./chroma_db4/data_mail4.csv')

def predict(query):
    res_output=[]
    persist_directory="./chroma_db4/"
    db2 = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
    db2.get(include=["embeddings","metadatas","documents"])
    docs = db2.similarity_search(query)
    #context=docs[0]
    context=docs[0].page_content
    question=query
    
    res_output.append(context)

    text=''
    text=str(data_emb['file'][docs[0].metadata['row']])
    res_output.append(text)
    
    nlp_ft = pipeline("question-answering", model="FlavioBF/multi-qa-mpnet-base-dot-v1_fine_tuned_model")
    text='ANSWER: '+str(nlp_ft(question=question,context=context)['answer'])
    text=text+'\n'
    text=text+'SCORE: '+str(nlp_ft(question=question,context=context)['score'])
    res_output.append(text)
    
    
    # Build prompt
    text =''
    template = """Use the context to answer the question at the end. If you don't know the answer, answer "I do not know the answer". Do not try to make up an answer. Use maximum two sentences maximum. \
    Keep the answer as concise as possible.
    {context}
    Question: {question}
    Answer:"""
    QA_CHAIN_PROMPT = PromptTemplate.from_template(template)
    # retrieve the context
    llm_name = "gpt-3.5-turbo"
    llm = ChatOpenAI(model_name=llm_name, temperature=0,openai_api_key=api_key)
    qa_chain = RetrievalQA.from_chain_type(
      llm,
      retriever=db2.as_retriever()
    )

    text=str(qa_chain({"query": question})['result'])
    res_output.append(text)

    # retrived context source storage id 
    text=''
    text='Context source: '
    text=text+str(data_emb['file'][int(docs[0].metadata['row'])])
    res_output.append(text)  
    return res_output

gr.Interface(
    fn=predict,

    inputs=[
        gr.Textbox(lines=2,label="Question"),
    ],
    examples=examples,
    title="Chat with Enron Dataset with RAG technique",
    description=description,
    outputs=[gr.Textbox(lines=8, label="CONTEXT"),
             gr.Textbox(lines=2,label="context source"),
            gr.Textbox(lines=3, label="multi-qa-mpnet-base-dot-v1_fine_tuned_model"),
             gr.Textbox(lines=4, label="GPT 3.5 turbo answer"),
             ]

).launch(share=True)