Jashu's picture
Update app.py
9a52fae verified
import os
from dotenv import load_dotenv
import gradio as gr
from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate, Settings
from llama_index.llms.huggingface import HuggingFaceInferenceAPI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from sentence_transformers import SentenceTransformer
load_dotenv()
# Configure the Llama index settings
Settings.llm = HuggingFaceInferenceAPI(
model_name="meta-llama/Meta-Llama-3-8B-Instruct",
tokenizer_name="meta-llama/Meta-Llama-3-8B-Instruct",
context_window=3000,
token=os.getenv("HF_TOKEN"),
max_new_tokens=512,
generate_kwargs={"temperature": 0.1},
)
Settings.embed_model = HuggingFaceEmbedding(
model_name="BAAI/bge-small-en-v1.5"
)
# Define the directory for persistent storage and data
PERSIST_DIR = "db"
PDF_DIRECTORY = 'data' # Changed to the directory containing PDFs
# Ensure directories exist
os.makedirs(PDF_DIRECTORY, exist_ok=True)
os.makedirs(PERSIST_DIR, exist_ok=True)
# Variable to store current chat conversation
current_chat_history = []
def data_ingestion_from_directory():
# Use SimpleDirectoryReader on the directory containing the PDF files
documents = SimpleDirectoryReader(PDF_DIRECTORY).load_data()
storage_context = StorageContext.from_defaults()
index = VectorStoreIndex.from_documents(documents)
index.storage_context.persist(persist_dir=PERSIST_DIR)
def handle_query(query):
chat_text_qa_msgs = [
(
"user",
"""
You are an expert on environmental science, specifically focused on air pollution. Your goal is to provide precise, informative, and actionable insights about air pollution, answering questions based on real-time and historical data on pollution levels, causes, and safety measures.
Your responses should be backed by recent research or data sources if available, and dont start from according from the document just give about the question related answer.
Context:
{context_str}
Question:
{query_str}
"""
)
]
text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
# Load index from storage
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
index = load_index_from_storage(storage_context)
# Use chat history to enhance response
context_str = ""
for past_query, response in reversed(current_chat_history):
if past_query.strip():
context_str += f"User asked: '{past_query}'\nBot answered: '{response}'\n"
query_engine = index.as_query_engine(text_qa_template=text_qa_template, context_str=context_str)
answer = query_engine.query(query)
if hasattr(answer, 'response'):
response = answer.response
elif isinstance(answer, dict) and 'response' in answer:
response = answer['response']
else:
response = "Sorry, as per my current knowledge I am unable to answer this question. Is there anything else I can help you with?"
# Remove sensitive information and unwanted sections from the response
sensitive_keywords = [PERSIST_DIR, PDF_DIRECTORY, "/", "\\", ".pdf", ".doc", ".txt"]
for keyword in sensitive_keywords:
response = response.replace(keyword, "")
# Remove sections starting with specific keywords
unwanted_sections = ["Page Label","Page Label:","page_label","page_label:","file_path:","file_path",]
for section in unwanted_sections:
if section in response:
response = response.split(section)[0]
# Additional cleanup for any remaining artifacts from replacements
response = ' '.join(response.split())
# Update current chat history
current_chat_history.append((query, response))
return response
# Example usage: Process PDF ingestion from directory
print("Processing PDF ingestion from directory:", PDF_DIRECTORY)
data_ingestion_from_directory()
# Define the input and output components for the Gradio interface
input_component = gr.Textbox(
show_label=False,
placeholder="Hey I Am Bujji, Ask me anything about Air Pollution..."
)
output_component = gr.Textbox()
# Function to handle queries
def chatbot_handler(query):
response = handle_query(query)
return response
# Create the Gradio interface
interface = gr.Interface(
fn=chatbot_handler,
inputs=input_component,
outputs=output_component,
title="Welcome to Air Pollution Solutions",
description="I am here to assist you with any questions you have about Pollution Preventions. How can I help you today?"
)
# Launch the Gradio interface
interface.launch()