Spaces:
Sleeping
Sleeping
from langchain_huggingface import HuggingFaceEmbeddings | |
from langchain_community.vectorstores import FAISS | |
from langchain_community.llms import CTransformers | |
from langchain.chains import RetrievalQA | |
import gradio as gr | |
from huggingface_hub import hf_hub_download | |
import os | |
DB_FAISS_PATH = "vectorstores/db_faiss" | |
def load_llm(): | |
model_name = 'TheBloke/Llama-2-7B-Chat-GGML' # Replace with the actual model repository name | |
model_path = hf_hub_download(repo_id=model_name, filename='pytorch_model.bin', cache_dir='./models') | |
llm = CTransformers( | |
model=model_path, | |
model_type="llama", | |
max_new_tokens=512, | |
temperature=0.5 | |
) | |
return llm | |
custom_prompt_template = """Use the following pieces of information to answer the user's question. | |
If you don't know the answer, just say that you don't know, don't try to make up an answer. | |
Context: {context} | |
Question: {question} | |
only return the helpful answer below and nothing else. | |
Helpful answer: | |
""" | |
def set_custom_prompt(): | |
prompt = PromptTemplate(template=custom_prompt_template, input_variables=['context', 'question']) | |
return prompt | |
def retrieval_QA_chain(llm, prompt, db): | |
qachain = RetrievalQA.from_chain_type( | |
llm=llm, | |
chain_type="stuff", | |
retriever=db.as_retriever(search_kwargs={'k': 2}), | |
return_source_documents=True, | |
chain_type_kwargs={'prompt': prompt} | |
) | |
return qachain | |
def qa_bot(): | |
embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-miniLM-L6-V2', model_kwargs={'device': 'cpu'}) | |
db = FAISS.load_local(DB_FAISS_PATH, embeddings, allow_dangerous_deserialization=True) | |
llm = load_llm() | |
qa_prompt = set_custom_prompt() | |
qa = retrieval_QA_chain(llm, qa_prompt, db) | |
return qa | |
def chatbot_response(query): | |
try: | |
qa = qa_bot() | |
response = qa({'query': query}) | |
answer = response["result"] | |
sources = response["source_documents"] | |
if sources: | |
answer += f"\nSources:" + str(sources) | |
else: | |
answer += "\nNo sources found" | |
return answer | |
except Exception as e: | |
return f"An error occurred: {str(e)}" | |
iface = gr.Interface( | |
fn=chatbot_response, | |
inputs=gr.Textbox(lines=2, placeholder="Enter your question..."), | |
outputs="text", | |
title="Medical Chatbot", | |
description="Ask a medical question and get answers based on the provided context.", | |
live=True | |
) | |
iface.launch() | |