from datasets import load_dataset
dataset = load_dataset("Namitg02/Test")
print(dataset)

from langchain.docstore.document import Document as LangchainDocument
from langchain.text_splitter import RecursiveCharacterTextSplitter
splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=15,separators=["\n\n", "\n", " ", ""])
docs = splitter.create_documents(str(dataset))


from langchain_community.embeddings import HuggingFaceEmbeddings
embedding_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

from langchain_community.vectorstores import Chroma
persist_directory = 'docs/chroma/'

vectordb = Chroma.from_documents(
    documents=docs,
    embedding=embedding_model,
    persist_directory=persist_directory
)



retriever = vectordb.as_retriever(
    search_type="similarity", search_kwargs={"k": 2}
)


from langchain.prompts import PromptTemplate
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory

memory = ConversationBufferMemory(
    memory_key="chat_history",
    return_messages=True
)

from transformers import pipeline
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
from langchain_core.messages import SystemMessage
from langchain_core.prompts import HumanMessagePromptTemplate
from langchain_core.prompts import ChatPromptTemplate
from langchain.prompts import PromptTemplate

print("check1")
question = "How can I reverse Diabetes?"


#template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Use three sentences maximum. Keep the answer as concise as possible. Always say "thanks for asking!" at the end of the answer. 
#{context}
#Question: {question}
#Helpful Answer:"""

#QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context", "question"],template=template)

from transformers import AutoTokenizer
from transformers import AutoModelForCausalLM

llm_model = "deepset/roberta-base-squad2"
tokenizer = AutoTokenizer.from_pretrained(llm_model)
model = AutoModelForCausalLM.from_pretrained(llm_model)
pipe = pipeline(model = llm_model, tokenizer = tokenizer,trust_remote_code=True, task = "question-answering", temperature=0.2)

question = "How can I reverse diabetes?"
#docs1 = retriever.invoke(question)
docs1 = retriever.invoke(get_relevant_documents)

print(docs1[0].page_content)
#print(docs1)[0]['generated_text'][-1]

print("check2")


#question = "How can I reverse diabetes?"
print("result")

print("check3")
#chain = pipe(question = question,context = "Use the following information to answer the question.  docs1[0].page_content.") 
chain = pipe(question = question,context = "Use the following information to answer the question. Diabetes can be cured by eating apples.") 


print("check3A")
print(chain)[0]['generated_text'][-1]
print("check3B")

import gradio as gr
ragdemo = gr.Interface.from_pipeline(chain)

print("check4")
ragdemo.launch()
print("check5")