File size: 4,412 Bytes
b47040f
 
 
 
71a31cb
 
 
b47040f
 
 
 
 
 
 
 
 
 
 
 
4ff5abb
b47040f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
71a31cb
b47040f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
71a31cb
b47040f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
71a31cb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import streamlit as st
import tempfile
import logging
from typing import List
from langchain.document_loaders import PyPDFLoader  # Updated import
from langchain.embeddings import HuggingFaceEmbeddings  # Updated import
from langchain.vectorstores import FAISS  # Updated import
from langchain.chains.summarize import load_summarize_chain
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import PromptTemplate
from transformers import pipeline

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Constants
EMBEDDING_MODEL = 'sentence-transformers/all-MiniLM-L6-v2'
DEFAULT_MODEL = "meta-llama/Meta-Llama-3.1-8B"

@st.cache_resource
def load_embeddings():
    """Load and cache the embedding model."""
    try:
        return HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
    except Exception as e:
        logger.error(f"Failed to load embeddings: {e}")
        st.error("Failed to load the embedding model. Please try again later.")
        return None

@st.cache_resource
def load_llm(model_name):
    """Load and cache the language model."""
    try:
        pipe = pipeline("text2text-generation", model=model_name, max_length=512)
        return pipe
    except Exception as e:
        logger.error(f"Failed to load LLM: {e}")
        st.error(f"Failed to load the model {model_name}. Please try again.")
        return None

def process_pdf(file) -> List[Document]:
    """Process the uploaded PDF file."""
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
            temp_file.write(file.getvalue())
            temp_file_path = temp_file.name
        
        loader = PyPDFLoader(file_path=temp_file_path)
        pages = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        documents = text_splitter.split_documents(pages)
        return documents
    except Exception as e:
        logger.error(f"Error processing PDF: {e}")
        st.error("Failed to process the PDF. Please make sure it's a valid PDF file.")
        return []

def create_vector_store(documents: List[Document], embeddings):
    """Create the vector store."""
    try:
        return FAISS.from_documents(documents, embeddings)
    except Exception as e:
        logger.error(f"Error creating vector store: {e}")
        st.error("Failed to create the vector store. Please try again.")
        return None

def summarize_report(documents: List[Document], llm) -> str:
    """Summarize the report using the loaded model."""
    try:
        prompt_template = """
        You are an AI specialized in summarizing comprehensive reports with a focus on funding, finances, and global comparisons. Given the detailed report content below, generate a concise and structured summary using bullet points and emojis...
        """

        prompt = PromptTemplate.from_template(prompt_template)
        chain = load_summarize_chain(llm, chain_type="stuff", prompt=prompt)
        summary = chain.invoke(documents)
        return summary['output_text']

    except Exception as e:
        logger.error(f"Error summarizing report: {e}")
        st.error("Failed to summarize the report. Please try again.")
        return ""

def main():
    st.title("Report Summarizer")
    
    model_option = st.sidebar.text_input("Enter model name", value=DEFAULT_MODEL)

    uploaded_file = st.sidebar.file_uploader("Upload your Report", type="pdf")

    llm = load_llm(model_option)
    embeddings = load_embeddings()

    if not llm or not embeddings:
        return

    if uploaded_file:
        with st.spinner("Processing PDF..."):
            documents = process_pdf(uploaded_file)

        if documents:
            with st.spinner("Creating vector store..."):
                db = create_vector_store(documents, embeddings)

            if db and st.button("Summarize"):
                with st.spinner(f"Generating structured summary using {model_option}..."):
                    summary = summarize_report(documents, llm)

                    if summary:
                        st.subheader("Structured Summary:")
                        st.markdown(summary)
                    else:
                        st.warning("Failed to generate summary. Please try again.")

if __name__ == "__main__":
    main()