Spaces:
Sleeping
Sleeping
import gradio as gr | |
from langchain.chains import RetrievalQA | |
from langchain.embeddings import HuggingFaceEmbeddings | |
from langchain.llms import HuggingFacePipeline | |
from langchain.memory import ConversationBufferWindowMemory | |
from langchain.vectorstores import Chroma | |
from langchain import PromptTemplate, LLMChain | |
from transformers import AutoTokenizer, pipeline | |
from typing import Dict, Any | |
import torch | |
# class AnswerConversationBufferMemory(ConversationBufferMemory): | |
class AnswerConversationBufferMemory(ConversationBufferWindowMemory): | |
def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None: | |
return super(AnswerConversationBufferMemory, self).save_context(inputs,{'response': outputs['result']}) | |
def clean_text(text): | |
# Remove excessive whitespace | |
cleaned_text = ' '.join(text.split()) | |
# Keep max one newline character | |
cleaned_text = cleaned_text.replace('\n\n', '\n') | |
return cleaned_text | |
def chatbot_llm_response(llm_response): | |
text = clean_text(llm_response['result']) + '\nSources:\n' | |
for source in llm_response["source_documents"]: | |
text += source.metadata['source'] + '\n' | |
return text | |
model_name = "databricks/dolly-v2-3b" | |
tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left") | |
generate_text = pipeline(model=model_name, | |
torch_dtype=torch.bfloat16, | |
trust_remote_code=True, | |
device_map="auto", | |
return_full_text=True, | |
max_new_tokens=256, | |
top_p=0.95, | |
top_k=50) | |
prompt = PromptTemplate( | |
input_variables=["instruction"], | |
template="{instruction}") | |
hf_pipeline = HuggingFacePipeline(pipeline=generate_text) | |
llm_chain = LLMChain(llm=hf_pipeline, prompt=prompt) | |
# top #2 when task = Retrieval June 2023 for under ~500 MB | |
model_name = "intfloat/e5-base-v2" | |
hf = HuggingFaceEmbeddings(model_name=model_name) | |
# Load up Vector Database | |
persist_directory = 'db' | |
vectordb = Chroma(persist_directory=persist_directory, | |
embedding_function=hf) | |
vectordb.get() | |
retriever = vectordb.as_retriever(search_kwargs={'k':3}) | |
# Configure Conversation Chain | |
memory = AnswerConversationBufferMemory(k=3) | |
qa_chain_with_memory = RetrievalQA.from_chain_type(llm=hf_pipeline, | |
chain_type="stuff", | |
retriever=retriever, | |
return_source_documents=True, | |
memory=memory) | |
# try to set the tone | |
template = ''' | |
You are the assistant to a tradesperson with knowledge of the Ontario Building Code. You provide specific details using the context given and the users question. | |
If you don't know the answer, you truthfully say you don't know and don't try to make up an answer. | |
---------------- | |
{context} | |
Question: {question} | |
Helpful Answer:''' | |
qa_chain_with_memory.combine_documents_chain.llm_chain.prompt.template = template | |
examples = ["What are the requirements for plumbing venting and drainage systems?", | |
"Summarize the electrical code regulation for wiring commercial buildings", | |
"Tell me the maximum allowable span for floor joists in residential construction", | |
"I'm looking for guidelines for fire assemblies and walls in tall buildings", | |
"What are the insulation requirements in new residential constructions?"] | |
def process_example(args): | |
for x in generate(args): | |
pass | |
return x | |
def generate(instruction): | |
response = qa_chain_with_memory(instruction) | |
processed_response = chatbot_llm_response(response) | |
result = "" | |
for word in processed_response.split(" "): | |
result += word + " " | |
yield result | |
with gr.Blocks(analytics_enabled=False) as demo: | |
with gr.Column(): | |
gr.Markdown("""# Dolly-Expert-Lite | |
Dolly-Expert-Lite is a bot for domain specific question | |
answering. Currently powered by the new Dolly-v2-3b open | |
source model. It's expert systems in the era of LLMs! | |
## Building Code Expert | |
In this example deployment, Dolly-Expert-Lite retrieves | |
information via a vector database made using the | |
[Ontario (Canada) Building Code](https://www.buildingcode.online) | |
sitemap LangChain loader. For details on the original Dolly | |
v2 model, please refer to the | |
[model card](https://huggingface.co/databricks/dolly-v2-12b) | |
### Type in the box below and click to ask the expert! | |
""" | |
) | |
with gr.Row(): | |
with gr.Column(scale=3): | |
instruction = gr.Textbox(placeholder="Enter your question here", label="Question", elem_id="q-input") | |
with gr.Box(): | |
gr.Markdown("**Answer**") | |
output = gr.Markdown(elem_id="q-output") | |
submit = gr.Button("Generate", variant="primary") | |
clear = gr.Button("Clear", variant="secondary") | |
gr.Examples( | |
examples=examples, | |
inputs=[instruction], | |
cache_examples=False, | |
fn=process_example, | |
outputs=[output], | |
) | |
submit.click(generate, inputs=[instruction], outputs=[output]) | |
clear.click(lambda: None, [], [output]) | |
instruction.submit(generate, inputs=[instruction], outputs=[output]) | |
demo.queue(concurrency_count=16).launch(debug=True) | |
demo.launch() |