|
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/username/repo/resolve/main/file1.pdf"] |
|
} |
|
|
|
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 i |
|
|
|
|
