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's the minimum pipe size needed for sinks, toilets, and showers?", | |
| "Are there any specific rules for installing backflow prevention devices?", | |
| "Can you guide me on the approved materials and methods for installing underground sewer lines?", | |
| "How much clearance is required for electrical panels and equipment like switchboards?", | |
| "Are there any restrictions or guidelines for outdoor electrical wiring and fixtures?", | |
| "Could you explain the proper bonding and grounding requirements for commercial buildings?", | |
| "What's the load-bearing capacity for beams and columns?", | |
| "Are there any specific rules for designing buildings to withstand earthquakes?", | |
| "Can you provide information on the fire resistance ratings for walls, floors, and roofs?", | |
| "What are the specific building code requirements for designing accessible entrances and pathways?", | |
| "Can you explain the regulations for fire protection systems and how they should be integrated into architectural designs?", | |
| "What are the foundation requirements in areas prone to earthquakes?", | |
| "Are there any restrictions or guidelines for installing electrical wiring and fixtures in wet locations?" | |
| ] | |
| 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() | |
| demo.launch() |