import streamlit as st
import os
import tempfile
from dotenv import load_dotenv
import PyPDF2
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.llms import HuggingFacePipeline
from langchain.chains import RetrievalQA
from langchain.vectorstores import FAISS
from langchain.document_loaders import PyPDFLoader
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch
# Load environment variables
load_dotenv()
# Configuration from environment variables
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
LLM_MODEL = os.getenv("LLM_MODEL", "microsoft/DialoGPT-medium")
CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "1000"))
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "200"))
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "256"))
TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7"))
STREAMLIT_PORT = int(os.getenv("STREAMLIT_SERVER_PORT", "8501"))
def validate_environment():
"""Validate environment variables and return status"""
status = {
"valid": True,
"warnings": [],
"errors": []
}
# Check if .env file exists
if not os.path.exists(".env"):
status["warnings"].append("No .env file found. Using default values.")
# Validate numeric values
if CHUNK_SIZE <= 0:
status["errors"].append("CHUNK_SIZE must be positive")
status["valid"] = False
if CHUNK_OVERLAP < 0:
status["errors"].append("CHUNK_OVERLAP must be non-negative")
status["valid"] = False
if MAX_TOKENS <= 0:
status["errors"].append("MAX_TOKENS must be positive")
status["valid"] = False
if not (0 <= TEMPERATURE <= 2):
status["warnings"].append("TEMPERATURE should be between 0 and 2")
return status
# Configure Streamlit page
st.set_page_config(
page_title="RAG PDF Chatbot",
page_icon="📚",
layout="wide",
initial_sidebar_state="expanded"
)
# Custom CSS for better UI
st.markdown("""
""", unsafe_allow_html=True)
class RAGChatbot:
def __init__(self):
self.embeddings = None
self.vectorstore = None
self.qa_chain = None
self.llm = None
self.documents = []
def initialize_models(self):
"""Initialize the embedding model and LLM"""
try:
# Initialize embeddings using sentence transformers (no API token needed)
self.embeddings = HuggingFaceEmbeddings(
model_name=EMBEDDING_MODEL,
model_kwargs={'device': 'cpu'},
encode_kwargs={'normalize_embeddings': True}
)
# Initialize LLM using local transformers (no API token needed)
tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
model = AutoModelForCausalLM.from_pretrained(
LLM_MODEL,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map="auto" if torch.cuda.is_available() else None
)
# Create pipeline
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=MAX_TOKENS,
temperature=TEMPERATURE,
do_sample=True,
pad_token_id=tokenizer.eos_token_id,
truncation=True
)
self.llm = HuggingFacePipeline(pipeline=pipe)
return True
except Exception as e:
st.error(f"Error initializing models: {str(e)}")
return False
def process_pdf(self, pdf_file):
"""Process uploaded PDF file and create vector store"""
try:
# Save uploaded file temporarily
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
tmp_file.write(pdf_file.read())
tmp_file_path = tmp_file.name
# Load PDF using PyPDFLoader
loader = PyPDFLoader(tmp_file_path)
self.documents = loader.load()
# Split documents into chunks
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
length_function=len,
)
texts = text_splitter.split_documents(self.documents)
# Create vector store
self.vectorstore = FAISS.from_documents(texts, self.embeddings)
# Create custom prompt template for better understanding
from langchain.prompts import PromptTemplate
custom_prompt = PromptTemplate(
template="""Context: {context}
Question: {question}
Answer:""",
input_variables=["context", "question"]
)
# Create QA chain with enhanced retrieval
self.qa_chain = RetrievalQA.from_chain_type(
llm=self.llm,
chain_type="stuff",
retriever=self.vectorstore.as_retriever(
search_type="similarity_score_threshold",
search_kwargs={
"k": 5, # Increased to get more relevant documents
"score_threshold": 0.3 # Lower threshold for more flexible matching
}
),
return_source_documents=True,
chain_type_kwargs={"prompt": custom_prompt}
)
# Clean up temporary file
os.unlink(tmp_file_path)
return True, len(texts)
except Exception as e:
st.error(f"Error processing PDF: {str(e)}")
return False, 0
def ask_question(self, question):
"""Ask a question and get answer from the RAG system"""
try:
if self.qa_chain is None:
return "Please upload a PDF file first.", []
# Truncate question if it's too long
if len(question) > 500:
question = question[:500] + "..."
# Use simple similarity search for better results
simple_retriever = self.vectorstore.as_retriever(search_kwargs={"k": 5})
docs = simple_retriever.get_relevant_documents(question)
if not docs:
return "I couldn't find relevant information in the document to answer your question. Please try rephrasing or asking about a different topic.", []
# Use the LLM directly with retrieved documents
context = "\n\n".join([doc.page_content for doc in docs])
# Try multiple prompt approaches for better results
prompts_to_try = [
f"""Here's some information: {context}
Question: {question}
Answer:""",
f"""Based on this text: {context}
{question}
Response:""",
f"""Context: {context}
Q: {question}
A:""",
f"""Information: {context}
{question}
Answer based on the information above:"""
]
answer = ""
for prompt in prompts_to_try:
try:
response = self.llm(prompt)
answer = response.strip()
# Check if we got a good answer
if (len(answer) > 20 and
answer.lower().strip() != question.lower().strip() and
not answer.lower().startswith(question.lower())):
break
except:
continue
# If still no good answer, use the last attempt
if not answer or len(answer) < 10:
answer = response.strip() if 'response' in locals() else ""
source_docs = docs
# Clean up the answer more aggressively
# Remove everything before the first meaningful content
lines = answer.split('\n')
cleaned_lines = []
found_content = False
for line in lines:
line = line.strip()
# Skip empty lines at the beginning
if not line and not found_content:
continue
# Skip template markers
if line.lower() in ['context:', 'question:', 'answer:']:
continue
# Skip lines that are just template text
if any(template_text in line.lower() for template_text in [
'use the following pieces',
'answer the following question',
'based on the provided context',
'according to the context',
'the context shows that',
'from the context'
]):
continue
# If we find actual content, start collecting
if line and not any(template_word in line.lower() for template_word in ['context:', 'question:', 'answer:']):
found_content = True
cleaned_lines.append(line)
answer = '\n'.join(cleaned_lines).strip()
# If answer is empty or too short, provide a fallback
if len(answer) < 10 or answer.lower().strip() == question.lower().strip():
# Provide a summary of the context as fallback
context_summary = context[:500] + "..." if len(context) > 500 else context
answer = f"Based on the document content, here's what I found: {context_summary}"
# If still no good answer, provide a generic response
if len(answer) < 20:
answer = "I found relevant information in the document, but I'm having trouble formatting the response. Please try rephrasing your question."
return answer, source_docs
except Exception as e:
error_msg = str(e)
if "max_length" in error_msg or "max_new_tokens" in error_msg:
return "The question or context is too long. Please try asking a shorter, more specific question.", []
else:
return f"Error getting answer: {error_msg}", []
def main():
# Header
st.markdown('PDF PARSER
', unsafe_allow_html=True)
st.markdown("Upload a PDF file and ask questions about its content!")
# Validate environment
env_status = validate_environment()
if not env_status["valid"]:
st.error("❌ Environment configuration errors:")
for error in env_status["errors"]:
st.error(f"• {error}")
return
if env_status["warnings"]:
for warning in env_status["warnings"]:
st.warning(f"⚠️ {warning}")
# Initialize session state
if "chatbot" not in st.session_state:
st.session_state.chatbot = RAGChatbot()
st.session_state.chat_history = []
st.session_state.pdf_processed = False
# Sidebar for file upload and settings
with st.sidebar:
st.markdown('", unsafe_allow_html=True)
# Main chat interface
if not st.session_state.pdf_processed:
st.info("👈 Please upload and process a PDF file using the sidebar to start chatting!")
# Show example questions
st.markdown("### 💡 Example Questions You Can Ask:")
example_questions = [
"What is the main topic of this document?",
"Can you summarize the key points?",
"What are the important findings or conclusions?",
"Are there any specific recommendations mentioned?",
"What methodology was used in this study?",
"Tell me about the results",
"What does this document say about...?",
"Explain the main concepts",
"What are the key takeaways?",
"How does this relate to...?"
]
for i, question in enumerate(example_questions, 1):
st.markdown(f"{i}. {question}")
st.info("💡 **Tip**: Ask questions in any way you like - the bot understands context and relevance!")
else:
# Chat interface
st.markdown("### 💬 Chat with your PDF")
# Display chat history
for message in st.session_state.chat_history:
if message["role"] == "user":
st.markdown(f"""
You: {message["content"]}
""", unsafe_allow_html=True)
else:
st.markdown(f"""
🤖 Assistant: {message["content"]}
""", unsafe_allow_html=True)
# Chat input
user_question = st.chat_input("Ask a question about the PDF content...")
if user_question:
# Add user message to chat history
st.session_state.chat_history.append({
"role": "user",
"content": user_question
})
# Get answer from chatbot
with st.spinner("Thinking..."):
answer, sources = st.session_state.chatbot.ask_question(user_question)
# Add bot response to chat history
st.session_state.chat_history.append({
"role": "assistant",
"content": answer,
"sources": sources
})
# Rerun to display new messages
st.rerun()
# Clear chat button
if st.button("🗑️ Clear Chat History"):
st.session_state.chat_history = []
st.rerun()
if __name__ == "__main__":
main()