StevanLS's picture
Nouvelle app
cf2de81
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.document_loaders import PyPDFLoader
from langchain import PromptTemplate
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferWindowMemory
import gradio as gr
from langchain_community.llms import Ollama
db = Chroma(persist_directory='content/Sugar cane/db', embedding_function=HuggingFaceEmbeddings())
retriever = db.as_retriever()
B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
instruction = "Given the context that has been provided. \n {context}, Answer the following question - \n{question}"
system_prompt = """You are an expert in sugar cane cultivation.
Be precise in your answers and give answers as concise as possible."""
def get_prompt(instruction, system_prompt):
SYSTEM_PROMPT = B_SYS + system_prompt + E_SYS
prompt_template = B_INST + SYSTEM_PROMPT + instruction + E_INST
return prompt_template
template = get_prompt(instruction, system_prompt)
prompt = PromptTemplate(template=template, input_variables=["context", "question"])
memory = ConversationBufferWindowMemory(
memory_key="chat_history", k=5,
return_messages=True
)
class ChatBot:
def __init__(self, memory, prompt, retriever = retriever):
self.memory = memory
self.prompt = prompt
self.retriever = retriever
def create_chat_bot(self):
llm = Ollama(model="llama2")
qa = ConversationalRetrievalChain.from_llm(
llm=llm,
retriever=self.retriever,
memory=self.memory,
combine_docs_chain_kwargs={"prompt": self.prompt}
)
return qa
chatbot = ChatBot(memory = memory, prompt = prompt)
bot = chatbot.create_chat_bot()
def clear_llm_memory():
bot.memory.clear()
def update_prompt(sys_prompt):
if sys_prompt == "":
sys_prompt = system_prompt
template = get_prompt(instruction, sys_prompt)
prompt = PromptTemplate(template=template, input_variables=["context", "question"])
bot.combine_docs_chain.llm_chain.prompt = prompt
with gr.Blocks() as demo:
chatbot = gr.Chatbot(label="Chat Bot", height = 300)
msg = gr.Textbox(label = "Question")
clear = gr.ClearButton([msg, chatbot])
clear_memory = gr.Button(value = "Clear LLM Memory")
def respond(message, chat_history):
bot_message = bot({"question": message})['answer']
chat_history.append((message, bot_message))
return "", chat_history
msg.submit(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
clear_memory.click(clear_llm_memory)
demo.launch(share=False, debug=True)