|
import os |
|
import requests |
|
import streamlit as st |
|
from io import BytesIO |
|
from PyPDF2 import PdfReader |
|
from langchain.text_splitter import RecursiveCharacterTextSplitter |
|
from langchain.embeddings import HuggingFaceEmbeddings |
|
from langchain.vectorstores import FAISS |
|
from transformers import pipeline |
|
import torch |
|
|
|
st.set_page_config(page_title="RAG-based PDF Chat", layout="centered", page_icon="📄") |
|
|
|
@st.cache_resource |
|
def load_summarization_pipeline(): |
|
try: |
|
summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=0 if torch.cuda.is_available() else -1) |
|
return summarizer |
|
except Exception as e: |
|
st.error(f"Failed to load the summarization model: {e}") |
|
return None |
|
|
|
summarizer = load_summarization_pipeline() |
|
|
|
PDF_FOLDERS = { |
|
"Folder 1": ["https://huggingface.co/spaces/ZeeAI1/LawTest3/tree/main/documents1"] |
|
} |
|
|
|
def fetch_pdf_text_from_folders(pdf_folders): |
|
all_text = "" |
|
for folder_name, urls in pdf_folders.items(): |
|
folder_text = f"\n[Folder: {folder_name}]\n" |
|
for url in urls: |
|
try: |
|
response = requests.get(url) |
|
response.raise_for_status() |
|
pdf_file = BytesIO(response.content) |
|
pdf_reader = PdfReader(pdf_file) |
|
for page in pdf_reader.pages: |
|
page_text = page.extract_text() |
|
if page_text: |
|
folder_text += page_text |
|
except Exception as e: |
|
st.error(f"Error fetching PDF from {url}: {e}") |
|
all_text += folder_text |
|
return all_text |
|
|
|
@st.cache_data |
|
def get_text_chunks(text): |
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200) |
|
return text_splitter.split_text(text) |
|
|
|
@st.cache_resource |
|
def load_embedding_function(): |
|
try: |
|
return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") |
|
except Exception as e: |
|
st.error(f"Failed to load embedding model: {e}") |
|
return None |
|
|
|
embedding_function = load_embedding_function() |
|
|
|
@st.cache_resource |
|
def load_or_create_vector_store(text_chunks): |
|
if not text_chunks: |
|
st.error("No valid text chunks found.") |
|
return None |
|
try: |
|
return FAISS.from_texts(text_chunks, embedding=embedding_function) |
|
except Exception as e: |
|
st.error(f"Failed to create or load vector store: {e}") |
|
return None |
|
|
|
def generate_summary_with_huggingface(query, retrieved_text): |
|
summarization_input = f"{query}\n\nRelated information:\n{retrieved_text}"[:1024] |
|
try: |
|
summary = summarizer(summarization_input, max_length=500, min_length=50, do_sample=False) |
|
return summary[0]["summary_text"] |
|
except Exception as e: |
|
st.error(f"Failed to generate summary: {e}") |
|
return "Error generating summary." |
|
|
|
def user_input(user_question, vector_store): |
|
if vector_store is None: |
|
return "Vector store is empty." |
|
try: |
|
docs = vector_store.similarity_search(user_question) |
|
context_text = " ".join([doc.page_content for doc in docs]) |
|
return generate_summary_with_huggingface(user_question, context_text) |
|
except Exception as e: |
|
st.error(f"Error in similarity search: {e}") |
|
return "Error in similarity search." |
|
|
|
def main(): |
|
st.title("📄 Gen AI Lawyers Guide") |
|
raw_text = fetch_pdf_text_from_folders(PDF_FOLDERS) |
|
text_chunks = get_text_chunks(raw_text) |
|
vector_store = load_or_create_vector_store(text_chunks) |
|
|
|
user_question = st.text_input("Ask a Question:", placeholder="Type your question here...") |
|
|
|
if st.button("Get Response"): |
|
if not user_question: |
|
st.warning("Please enter a question before submitting.") |
|
else: |
|
with st.spinner("Generating response..."): |
|
answer = user_input(user_question, vector_store) |
|
st.markdown(f"**🤖 AI:** {answer}") |
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
|
|
