rohitashva's picture
Update app.py
f6c2e5c verified
import streamlit as st
import logging
import os
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import CSVLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain.llms import HuggingFaceHub
import dotenv
import yaml
import os
import zipfile
zip_file = "faiss_index.zip"
with zipfile.ZipFile(zip_file, 'r') as zip_ref:
zip_ref.extractall(".") # Extract to the current directory
print("Unzipping completed successfully.")
dotenv.load_dotenv()
def load_config():
with open("yaml-editor-online.yaml", "r") as f:
config = yaml.safe_load(f)
return config
hf_token = os.getenv("HUGGING")
config = load_config()
logging.basicConfig(level=logging.INFO)
embeddings_model = HuggingFaceEmbeddings(model_name=config["embedding_model"])
def create_vector_db():
try:
loader = CSVLoader(file_path="disease.csv", source_column="Disease Information")
data = loader.load()
vectordb = FAISS.from_documents(documents=data, embedding=embeddings_model)
vectordb.save_local(config["vector_db_path"])
logging.info("Vector database successfully created and saved.")
except Exception as e:
logging.error("Error creating vector database:", exc_info=e)
def get_qa_chain(query):
try:
if not os.path.exists(config["vector_db_path"]):
logging.error(f"FAISS index path does not exist: {config['vector_db_path']}")
return "Error: No data found."
vectordb = FAISS.load_local(
config["vector_db_path"], embeddings_model, allow_dangerous_deserialization=True
)
retriever = vectordb.as_retriever(score_threshold=config["score_threshold"])
relevant_docs = retriever.get_relevant_documents(query)[:3]
if not relevant_docs:
return "No relevant information found."
summarized_context = " ".join(doc.page_content for doc in relevant_docs)
prompt_template = """
Given the following health-related context and a question, generate a structured answer:
QUESTION: {query}
Ensure the response is easy to understand and medically accurate.
"""
prompt = PromptTemplate(input_variables=["query"], template=prompt_template).format(query=query)
llm = HuggingFaceHub(
repo_id=config["model_name"],
model_kwargs={
"temperature": config["temperature"],
"max_length": 200,
"top_p": config["top_p"],
"top_k": config["top_k"]
},
huggingfacehub_api_token=hf_token
)
response = llm(prompt)
return response.strip()
except Exception as e:
logging.error("Error getting response:", exc_info=e)
return "Sorry, there was an error processing your request."
def main():
st.set_page_config(page_title="Health Disease Chatbot", page_icon="🩺", layout="centered")
st.markdown(
"""
<style>
.stApp {
background-color: #f0f2f6;
color: #333;
font-family: 'Arial', sans-serif;
}
.title {
color: #2E7D32;
text-align: center;
}
.query-input {
border-radius: 10px;
padding: 10px;
}
.response-box {
background-color: #ffffff;
padding: 15px;
border-radius: 8px;
box-shadow: 2px 2px 10px rgba(0,0,0,0.1);
}
</style>
""",
unsafe_allow_html=True
)
st.markdown("<h1 class='title'>🩺 Health Disease Chatbot</h1>", unsafe_allow_html=True)
st.write("Enter a question related to health conditions, symptoms, or treatments.")
query = st.text_input("Your health-related question:", key="query", help="Ask about diseases, symptoms, or treatments.")
if st.button("Get Information"):
if query:
response = get_qa_chain(query)
st.markdown(f"<div class='response-box'><b>Response:</b><br>{response}</div>", unsafe_allow_html=True)
else:
st.warning("Please enter a query to get a response.")
if __name__ == "__main__":
if not os.path.exists(config["vector_db_path"]):
logging.info(f"Vector database not found at {config['vector_db_path']}, creating it now.")
create_vector_db()
main()